{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatCompletion(id='chatcmpl-9epZmBcCQkmy3kvxBBp8aFoaMOROJ', choices=[Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, role='assistant', function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_3LeBAkryCJcDQ6ZQKkKJUJgr', function=Function(arguments='{\"location\": \"San Francisco, CA\"}', name='get_current_weather'), type='function'), ChatCompletionMessageToolCall(id='call_cZublzUU6nLPiY2sa8M6vg0L', function=Function(arguments='{\"location\": \"Tokyo, Japan\"}', name='get_current_weather'), type='function'), ChatCompletionMessageToolCall(id='call_mLffU4B69qExzM8qaeyzkYQd', function=Function(arguments='{\"location\": \"Paris, France\"}', name='get_current_weather'), type='function')]))], created=1719518406, model='gpt-4o-2024-05-13', object='chat.completion', service_tier=None, system_fingerprint='fp_ce0793330f', usage=CompletionUsage(completion_tokens=68, prompt_tokens=85, total_tokens=153))\n",
      "ChatCompletion(id='chatcmpl-9epZo59ilaWhx9uHAV7cWZG8rWrjb', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content=\"Currently, here's the weather in the specified cities:\\n\\n- **San Francisco, CA:** 72°F\\n- **Tokyo, Japan:** 10°C\\n- **Paris, France:** 22°C\", role='assistant', function_call=None, tool_calls=None))], created=1719518408, model='gpt-4o-2024-05-13', object='chat.completion', service_tier=None, system_fingerprint='fp_d576307f90', usage=CompletionUsage(completion_tokens=40, prompt_tokens=151, total_tokens=191))\n"
     ]
    }
   ],
   "source": [
    "from openai import OpenAI\n",
    "import json\n",
    "\n",
    "client = OpenAI()\n",
    "\n",
    "\n",
    "# Example dummy function hard coded to return the same weather\n",
    "# In production, this could be your backend API or an external API\n",
    "def get_current_weather(location, unit=\"fahrenheit\"):\n",
    "    \"\"\"Get the current weather in a given location\"\"\"\n",
    "    if \"tokyo\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Tokyo\", \"temperature\": \"10\", \"unit\": unit})\n",
    "    elif \"san francisco\" in location.lower():\n",
    "        return json.dumps(\n",
    "            {\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": unit}\n",
    "        )\n",
    "    elif \"paris\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": unit})\n",
    "    else:\n",
    "        return json.dumps({\"location\": location, \"temperature\": \"unknown\"})\n",
    "\n",
    "\n",
    "def run_conversation():\n",
    "    # Step 1: send the conversation and available functions to the model\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"What's the weather like in San Francisco, Tokyo, and Paris?\",\n",
    "        }\n",
    "    ]\n",
    "    tools = [\n",
    "        {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": \"get_current_weather\",\n",
    "                \"description\": \"Get the current weather in a given location\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"location\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The city and state, e.g. San Francisco, CA\",\n",
    "                        },\n",
    "                        \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"]},\n",
    "                    },\n",
    "                    \"required\": [\"location\"],\n",
    "                },\n",
    "            },\n",
    "        }\n",
    "    ]\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        messages=messages,\n",
    "        tools=tools,\n",
    "        tool_choice=\"auto\",  # auto is default, but we'll be explicit\n",
    "    )\n",
    "    print(response)\n",
    "    response_message = response.choices[0].message\n",
    "    tool_calls = response_message.tool_calls\n",
    "    # Step 2: check if the model wanted to call a function\n",
    "    if tool_calls:\n",
    "        # Step 3: call the function\n",
    "        # Note: the JSON response may not always be valid; be sure to handle errors\n",
    "        available_functions = {\n",
    "            \"get_current_weather\": get_current_weather,\n",
    "        }  # only one function in this example, but you can have multiple\n",
    "        messages.append(response_message)  # extend conversation with assistant's reply\n",
    "        # Step 4: send the info for each function call and function response to the model\n",
    "        for tool_call in tool_calls:\n",
    "            function_name = tool_call.function.name\n",
    "            function_to_call = available_functions[function_name]\n",
    "            function_args = json.loads(tool_call.function.arguments)\n",
    "            function_response = function_to_call(\n",
    "                location=function_args.get(\"location\"),\n",
    "                unit=function_args.get(\"unit\"),\n",
    "            )\n",
    "            messages.append(\n",
    "                {\n",
    "                    \"tool_call_id\": tool_call.id,\n",
    "                    \"role\": \"tool\",\n",
    "                    \"name\": function_name,\n",
    "                    \"content\": function_response,\n",
    "                }\n",
    "            )  # extend conversation with function response\n",
    "        second_response = client.chat.completions.create(\n",
    "            model=\"gpt-4o\",\n",
    "            messages=messages,\n",
    "        )  # get a new response from the model where it can see the function response\n",
    "        return second_response\n",
    "\n",
    "\n",
    "print(run_conversation())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from typing import Any, Dict, List\n",
    "import numpy as np\n",
    "import time\n",
    "import asyncio\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two numbers.\"\"\"\n",
    "    time.sleep(1)\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two numbers.\"\"\"\n",
    "    time.sleep(1)\n",
    "    return a + b\n",
    "\n",
    "\n",
    "async def divide(a: float, b: float) -> float:\n",
    "    \"\"\"Divide two numbers.\"\"\"\n",
    "    await asyncio.sleep(1)\n",
    "    return float(a) / b\n",
    "\n",
    "\n",
    "async def search(query: str) -> List[str]:\n",
    "    \"\"\"Search for query and return a list of results.\"\"\"\n",
    "    await asyncio.sleep(1)\n",
    "    return [\"result1\" + query, \"result2\" + query]\n",
    "\n",
    "\n",
    "def numpy_sum(arr: np.ndarray) -> float:\n",
    "    \"\"\"Sum the elements of an array.\"\"\"\n",
    "    return np.sum(arr)\n",
    "\n",
    "\n",
    "x = 2\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Point:\n",
    "    x: int\n",
    "    y: int\n",
    "\n",
    "\n",
    "def add_points(p1: Point, p2: Point) -> Point:\n",
    "    return Point(p1.x + p2.x, p1.y + p2.y)\n",
    "\n",
    "\n",
    "all_functions = [multiply, add, divide, search, numpy_sum, add_points]\n",
    "\n",
    "all_functions_dict = {f.__name__: f for f in all_functions}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FunctionTool(fn: <function multiply at 0x127bffb00>, async: False, definition: FunctionDefinition(func_name='multiply', func_desc='multiply(a: int, b: int) -> int\\nMultiply two numbers.', func_parameters={'type': 'object', 'properties': {'a': {'type': 'int'}, 'b': {'type': 'int'}}, 'required': ['a', 'b']}))\n",
      "FunctionTool(fn: <function add at 0x127bff420>, async: False, definition: FunctionDefinition(func_name='add', func_desc='add(a: int, b: int) -> int\\nAdd two numbers.', func_parameters={'type': 'object', 'properties': {'a': {'type': 'int'}, 'b': {'type': 'int'}}, 'required': ['a', 'b']}))\n",
      "FunctionTool(fn: <function divide at 0x127bffd80>, async: True, definition: FunctionDefinition(func_name='divide', func_desc='divide(a: float, b: float) -> float\\nDivide two numbers.', func_parameters={'type': 'object', 'properties': {'a': {'type': 'float'}, 'b': {'type': 'float'}}, 'required': ['a', 'b']}))\n",
      "FunctionTool(fn: <function search at 0x127bff920>, async: True, definition: FunctionDefinition(func_name='search', func_desc='search(query: str) -> List[str]\\nSearch for query and return a list of results.', func_parameters={'type': 'object', 'properties': {'query': {'type': 'str'}}, 'required': ['query']}))\n",
      "FunctionTool(fn: <function numpy_sum at 0x127bfdd00>, async: False, definition: FunctionDefinition(func_name='numpy_sum', func_desc='numpy_sum(arr: numpy.ndarray) -> float\\nSum the elements of an array.', func_parameters={'type': 'object', 'properties': {'arr': {'type': 'ndarray'}}, 'required': ['arr']}))\n",
      "FunctionTool(fn: <function add_points at 0x127bff4c0>, async: False, definition: FunctionDefinition(func_name='add_points', func_desc='add_points(p1: __main__.Point, p2: __main__.Point) -> __main__.Point\\nNone', func_parameters={'type': 'object', 'properties': {'p1': {'type': 'Point', 'properties': {'x': {'type': 'int'}, 'y': {'type': 'int'}}, 'required': ['x', 'y']}, 'p2': {'type': 'Point', 'properties': {'x': {'type': 'int'}, 'y': {'type': 'int'}}, 'required': ['x', 'y']}}, 'required': ['p1', 'p2']}))\n"
     ]
    }
   ],
   "source": [
    "# describing the functions\n",
    "\n",
    "from adalflow.core.func_tool import FunctionTool\n",
    "\n",
    "functions = [multiply, add, divide, search, numpy_sum, add_points]\n",
    "tools = [FunctionTool(fn=fn) for fn in functions]\n",
    "for tool in tools:\n",
    "    print(tool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a context map\n",
    "context_map = {tool.definition.func_name: tool for tool in tools}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'func_name': 'numpy_sum', 'func_desc': 'numpy_sum(arr: numpy.ndarray) -> float\\nSum the elements of an array.', 'func_parameters': {'type': 'object', 'properties': {'arr': {'type': 'ndarray'}}, 'required': ['arr']}}\n",
      "{\n",
      "    \"func_name\": \"numpy_sum\",\n",
      "    \"func_desc\": \"numpy_sum(arr: numpy.ndarray) -> float\\nSum the elements of an array.\",\n",
      "    \"func_parameters\": {\n",
      "        \"type\": \"object\",\n",
      "        \"properties\": {\n",
      "            \"arr\": {\n",
      "                \"type\": \"ndarray\"\n",
      "            }\n",
      "        },\n",
      "        \"required\": [\n",
      "            \"arr\"\n",
      "        ]\n",
      "    }\n",
      "}\n",
      "\"func_name: numpy_sum\\nfunc_desc: 'numpy_sum(arr: numpy.ndarray) -> float\\n\\n  Sum the elements of an array.'\\nfunc_parameters:\\n  type: object\\n  properties:\\n    arr:\\n      type: ndarray\\n  required:\\n  - arr\\n\"\n"
     ]
    }
   ],
   "source": [
    "print(tools[-2].definition.to_dict())\n",
    "\n",
    "print(tools[-2].definition.to_json())\n",
    "\n",
    "print(repr(tools[-2].definition.to_yaml()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'func_name': 'get_current_weather',\n",
       " 'func_desc': \"get_current_weather(location, unit='fahrenheit')\\nGet the current weather in a given location\",\n",
       " 'func_parameters': {'type': 'object',\n",
       "  'properties': {'location': {'type': 'Any'},\n",
       "   'unit': {'type': 'Any', 'default': 'fahrenheit'}},\n",
       "  'required': ['location']}}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# tool definition for get_current_weather\n",
    "\n",
    "ft = FunctionTool(fn=get_current_weather)\n",
    "ft.definition.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ToolManager(Tools: [FunctionTool(fn: <function multiply at 0x127bffb00>, async: False, definition: FunctionDefinition(func_name='multiply', func_desc='multiply(a: int, b: int) -> int\\nMultiply two numbers.', func_parameters={'type': 'object', 'properties': {'a': {'type': 'int'}, 'b': {'type': 'int'}}, 'required': ['a', 'b']})), FunctionTool(fn: <function add at 0x127bff420>, async: False, definition: FunctionDefinition(func_name='add', func_desc='add(a: int, b: int) -> int\\nAdd two numbers.', func_parameters={'type': 'object', 'properties': {'a': {'type': 'int'}, 'b': {'type': 'int'}}, 'required': ['a', 'b']})), FunctionTool(fn: <function divide at 0x127bffd80>, async: True, definition: FunctionDefinition(func_name='divide', func_desc='divide(a: float, b: float) -> float\\nDivide two numbers.', func_parameters={'type': 'object', 'properties': {'a': {'type': 'float'}, 'b': {'type': 'float'}}, 'required': ['a', 'b']})), FunctionTool(fn: <function search at 0x127bff920>, async: True, definition: FunctionDefinition(func_name='search', func_desc='search(query: str) -> List[str]\\nSearch for query and return a list of results.', func_parameters={'type': 'object', 'properties': {'query': {'type': 'str'}}, 'required': ['query']})), FunctionTool(fn: <function numpy_sum at 0x127bfdd00>, async: False, definition: FunctionDefinition(func_name='numpy_sum', func_desc='numpy_sum(arr: numpy.ndarray) -> float\\nSum the elements of an array.', func_parameters={'type': 'object', 'properties': {'arr': {'type': 'ndarray'}}, 'required': ['arr']})), FunctionTool(fn: <function add_points at 0x127bff4c0>, async: False, definition: FunctionDefinition(func_name='add_points', func_desc='add_points(p1: __main__.Point, p2: __main__.Point) -> __main__.Point\\nNone', func_parameters={'type': 'object', 'properties': {'p1': {'type': 'Point', 'properties': {'x': {'type': 'int'}, 'y': {'type': 'int'}}, 'required': ['x', 'y']}, 'p2': {'type': 'Point', 'properties': {'x': {'type': 'int'}, 'y': {'type': 'int'}}, 'required': ['x', 'y']}}, 'required': ['p1', 'p2']}))], Additional Context: {})\n"
     ]
    }
   ],
   "source": [
    "# to further help us manage the whole process, we will use a tool manager\n",
    "\n",
    "from adalflow.core.tool_manager import ToolManager\n",
    "\n",
    "tool_manager = ToolManager(tools=functions)\n",
    "print(tool_manager)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionOutput(name='get_current_weather', input=Function(thought=None, name='get_current_weather', args=(), kwargs={'location': 'San Francisco', 'unit': 'celsius'}), parsed_input=None, output='{\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": \"celsius\"}', error=None)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# execute get_current_weather using function call\n",
    "\n",
    "ft.call(**{\"location\": \"San Francisco\", \"unit\": \"celsius\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Task pending name='Task-1' coro=<FunctionTool.acall() running at /Users/liyin/Documents/test/LightRAG/lightrag/core/func_tool.py:124>>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FunctionOutput(name='divide', input=Function(thought=None, name='divide', args=(), kwargs={'a': 10, 'b': 2}), parsed_input=None, output=5.0, error=None)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "FunctionOutput(name='divide', input=Function(thought=None, name='divide', args=(), kwargs={'a': 10, 'b': 2}), parsed_input=None, output=5.0, error=None)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# async call\n",
    "import nest_asyncio\n",
    "from IPython.display import display\n",
    "\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "# call it synchronously using execute\n",
    "\n",
    "print(tools[2].execute(**{\"a\": 10, \"b\": 2}))\n",
    "\n",
    "display(await tools[2].acall(**{\"a\": 10, \"b\": 2}))\n",
    "display(await tools[2].execute(**{\"a\": 10, \"b\": 2}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<coroutine object to_thread at 0x127be1a40>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/h8/nhgbdr4d18x2r49j4pk5z6gw0000gn/T/ipykernel_1588/1960468447.py:5: RuntimeWarning: coroutine 'to_thread' was never awaited\n",
      "  print(tools[1].execute(**{\"a\": 10, \"b\": 2}))\n",
      "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FunctionOutput(name='add', input=Function(thought=None, name='add', args=(), kwargs={'a': 10, 'b': 2}), parsed_input=None, output=12, error=None)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FunctionOutput(name='add', input=Function(thought=None, name='add', args=(), kwargs={'a': 10, 'b': 2}), parsed_input=None, output=12, error=None)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# run sync func\n",
    "\n",
    "# in sync way\n",
    "\n",
    "print(tools[1].execute(**{\"a\": 10, \"b\": 2}))\n",
    "print(tools[1].call(**{\"a\": 10, \"b\": 2}))\n",
    "\n",
    "# in async way\n",
    "\n",
    "display(await tools[1].execute(**{\"a\": 10, \"b\": 2}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "run_sync_and_async_mix_without_wait time: 4.006277084350586\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<Task pending name='Task-4' coro=<FunctionTool.acall() running at /Users/liyin/Documents/test/LightRAG/lightrag/core/func_tool.py:124>>,\n",
       " FunctionOutput(name='sync_function_2', input=Function(thought=None, name='sync_function_2', args=(), kwargs={}), parsed_input=None, output='Function 2 completed', error=None),\n",
       " FunctionOutput(name='sync_function_2', input=Function(thought=None, name='sync_function_2', args=(), kwargs={}), parsed_input=None, output='Function 2 completed', error=None)]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "run_sync_and_async_mix time: 2.0017499923706055\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[FunctionOutput(name='async_function_1', input=Function(thought=None, name='async_function_1', args=(), kwargs={}), parsed_input=None, output='Function 1 completed', error=None),\n",
       " FunctionOutput(name='sync_function_2', input=Function(thought=None, name='sync_function_2', args=(), kwargs={}), parsed_input=None, output='Function 2 completed', error=None),\n",
       " FunctionOutput(name='async_function_2', input=Function(thought=None, name='async_function_2', args=(), kwargs={}), parsed_input=None, output='Function 2 completed', error=None)]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# call all the above functions\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "\n",
    "async def async_function_1():\n",
    "    await asyncio.sleep(1)\n",
    "    return \"Function 1 completed\"\n",
    "\n",
    "\n",
    "def sync_function_1():\n",
    "    time.sleep(1)\n",
    "    return \"Function 1 completed\"\n",
    "\n",
    "\n",
    "async def async_function_2():\n",
    "    await asyncio.sleep(2)\n",
    "    return \"Function 2 completed\"\n",
    "\n",
    "\n",
    "def sync_function_2():\n",
    "    time.sleep(2)\n",
    "    return \"Function 2 completed\"\n",
    "\n",
    "\n",
    "async_tool_1 = FunctionTool(async_function_1)\n",
    "sync_tool_1 = FunctionTool(sync_function_2)\n",
    "async_tool_2 = FunctionTool(async_function_2)\n",
    "sync_tool_2 = FunctionTool(sync_function_2)\n",
    "\n",
    "\n",
    "def run_sync_and_async_mix_without_wait():\n",
    "    # both sync and async tool can use execute\n",
    "    # sync tool can also use call\n",
    "    # takes 5 seconds (1+1+2) + overhead\n",
    "    start_time = time.time()\n",
    "    results = [\n",
    "        async_tool_1.execute(),\n",
    "        sync_tool_1.call(),\n",
    "        sync_tool_2.call(),\n",
    "    ]\n",
    "    end_time = time.time()\n",
    "    print(f\"run_sync_and_async_mix_without_wait time: {end_time - start_time}\")\n",
    "    return results\n",
    "\n",
    "\n",
    "async def run_sync_and_async_mix():\n",
    "    # both sync and async tool can use execute&to_thread\n",
    "    # async tool can also use acall without to_thread\n",
    "    # takes a bit over 2 seconds max(2)\n",
    "    start_time = time.time()\n",
    "    results = await asyncio.gather(\n",
    "        async_tool_1.execute(),\n",
    "        sync_tool_1.execute(),\n",
    "        async_tool_2.acall(),\n",
    "    )\n",
    "    end_time = time.time()\n",
    "    print(f\"run_sync_and_async_mix time: {end_time - start_time}\")\n",
    "    return results\n",
    "\n",
    "\n",
    "# Execute functions\n",
    "results_without_wait = run_sync_and_async_mix_without_wait()\n",
    "display(results_without_wait)\n",
    "\n",
    "results_with_wait = asyncio.run(run_sync_and_async_mix())\n",
    "display(results_with_wait)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare a template for generator\n",
    "template = r\"\"\"<SYS>You have these tools available:\n",
    "{% if tools %}\n",
    "<TOOLS>\n",
    "{% for tool in tools %}\n",
    "{{ loop.index }}.\n",
    "{{tool}}\n",
    "------------------------\n",
    "{% endfor %}\n",
    "</TOOLS>\n",
    "{% endif %}\n",
    "<OUTPUT_FORMAT>\n",
    "{{output_format_str}}\n",
    "</OUTPUT_FORMAT>\n",
    "</SYS>\n",
    "User: {{input_str}}\n",
    "You:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<SYS>You have these tools available:\n",
      "<TOOLS>\n",
      "1.\n",
      "func_name: multiply\n",
      "func_desc: 'multiply(a: int, b: int) -> int\n",
      "\n",
      "  Multiply two numbers.'\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: int\n",
      "    b:\n",
      "      type: int\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "2.\n",
      "func_name: add\n",
      "func_desc: 'add(a: int, b: int) -> int\n",
      "\n",
      "  Add two numbers.'\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: int\n",
      "    b:\n",
      "      type: int\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "3.\n",
      "func_name: divide\n",
      "func_desc: 'divide(a: float, b: float) -> float\n",
      "\n",
      "  Divide two numbers.'\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: float\n",
      "    b:\n",
      "      type: float\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "4.\n",
      "func_name: search\n",
      "func_desc: 'search(query: str) -> List[str]\n",
      "\n",
      "  Search for query and return a list of results.'\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    query:\n",
      "      type: str\n",
      "  required:\n",
      "  - query\n",
      "\n",
      "------------------------\n",
      "5.\n",
      "func_name: numpy_sum\n",
      "func_desc: 'numpy_sum(arr: numpy.ndarray) -> float\n",
      "\n",
      "  Sum the elements of an array.'\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    arr:\n",
      "      type: ndarray\n",
      "  required:\n",
      "  - arr\n",
      "\n",
      "------------------------\n",
      "6.\n",
      "func_name: add_points\n",
      "func_desc: 'add_points(p1: __main__.Point, p2: __main__.Point) -> __main__.Point\n",
      "\n",
      "  None'\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    p1:\n",
      "      type: Point\n",
      "      properties:\n",
      "        x:\n",
      "          type: int\n",
      "        y:\n",
      "          type: int\n",
      "      required:\n",
      "      - x\n",
      "      - y\n",
      "    p2:\n",
      "      type: Point\n",
      "      properties:\n",
      "        x:\n",
      "          type: int\n",
      "        y:\n",
      "          type: int\n",
      "      required:\n",
      "      - x\n",
      "      - y\n",
      "  required:\n",
      "  - p1\n",
      "  - p2\n",
      "\n",
      "------------------------\n",
      "</TOOLS>\n",
      "<OUTPUT_FORMAT>\n",
      "None\n",
      "</OUTPUT_FORMAT>\n",
      "</SYS>\n",
      "User: None\n",
      "You:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# let's see how the template can be rendered with tools\n",
    "from adalflow.core.prompt_builder import Prompt\n",
    "\n",
    "prompt = Prompt(template=template)\n",
    "small_tool_manager = ToolManager(tools=tools[:2])\n",
    "\n",
    "renered_prompt = prompt(tools=tool_manager.yaml_definitions)\n",
    "print(renered_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<SYS>You have these tools available:\n",
      "<OUTPUT_FORMAT>\n",
      "{\n",
      "    \"name\": \"The name of the function (str) (optional)\",\n",
      "    \"args\": \"The positional arguments of the function (Optional) (optional)\",\n",
      "    \"kwargs\": \"The keyword arguments of the function (Optional) (optional)\"\n",
      "}\n",
      "</OUTPUT_FORMAT>\n",
      "</SYS>\n",
      "User: None\n",
      "You:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# let's render the output format using Function class\n",
    "\n",
    "from adalflow.core.types import Function\n",
    "\n",
    "\n",
    "output_data_class = Function\n",
    "output_format_str = output_data_class.to_json_signature(exclude=[\"thought\"])\n",
    "\n",
    "renered_prompt = prompt(output_format_str=output_format_str)\n",
    "print(renered_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<SYS>You have these tools available:\n",
      "<OUTPUT_FORMAT>\n",
      "{\n",
      "    \"action\": \"Formatted as FuncName(<args>, <kwargs>), where FuncName is the function name, <args> are positional arguments, and <kwargs> are keyword arguments in key=value form. Example: 'FuncName(1, b=2)' calls 'FuncName' with positional argument 1 and keyword argument b=2. (str) (required)\"\n",
      "}\n",
      "</OUTPUT_FORMAT>\n",
      "</SYS>\n",
      "User: None\n",
      "You:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# use FunctionExpression\n",
    "from adalflow.core.types import FunctionExpression\n",
    "\n",
    "output_data_class = FunctionExpression\n",
    "output_format_str = output_data_class.to_json_signature(exclude=[\"thought\"])\n",
    "print(prompt(output_format_str=output_format_str))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your output should be formatted as a standard JSON instance with the following schema:\n",
      "```\n",
      "{\n",
      "    \"name\": \"The name of the function (str) (optional)\",\n",
      "    \"args\": \"The positional arguments of the function (Optional) (optional)\",\n",
      "    \"kwargs\": \"The keyword arguments of the function (Optional) (optional)\"\n",
      "}\n",
      "```\n",
      "-Make sure to always enclose the JSON output in triple backticks (```). Please do not add anything other than valid JSON output!\n",
      "-Use double quotes for the keys and string values.\n",
      "-Follow the JSON formatting conventions.\n"
     ]
    }
   ],
   "source": [
    "# let's adds more instruction and this time, we will use JsonOutputParser\n",
    "\n",
    "from adalflow.components.output_parsers import JsonOutputParser\n",
    "\n",
    "func_parser = JsonOutputParser(data_class=Function)\n",
    "instructions = func_parser.format_instructions(exclude=[\"thought\"])\n",
    "print(instructions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Generator(\n",
       "  model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
       "  (prompt): Prompt(\n",
       "    template: <SYS>You have these tools available:\n",
       "    {% if tools %}\n",
       "    <TOOLS>\n",
       "    {% for tool in tools %}\n",
       "    {{ loop.index }}.\n",
       "    {{tool}}\n",
       "    ------------------------\n",
       "    {% endfor %}\n",
       "    </TOOLS>\n",
       "    {% endif %}\n",
       "    <OUTPUT_FORMAT>\n",
       "    {{output_format_str}}\n",
       "    </OUTPUT_FORMAT>\n",
       "    </SYS>\n",
       "    User: {{input_str}}\n",
       "    You:\n",
       "    , prompt_kwargs: {'tools': [\"func_name: multiply\\nfunc_desc: 'multiply(a: int, b: int) -> int\\n\\n  Multiply two numbers.'\\nfunc_parameters:\\n  type: object\\n  properties:\\n    a:\\n      type: int\\n    b:\\n      type: int\\n  required:\\n  - a\\n  - b\\n\", \"func_name: add\\nfunc_desc: 'add(a: int, b: int) -> int\\n\\n  Add two numbers.'\\nfunc_parameters:\\n  type: object\\n  properties:\\n    a:\\n      type: int\\n    b:\\n      type: int\\n  required:\\n  - a\\n  - b\\n\", \"func_name: divide\\nfunc_desc: 'divide(a: float, b: float) -> float\\n\\n  Divide two numbers.'\\nfunc_parameters:\\n  type: object\\n  properties:\\n    a:\\n      type: float\\n    b:\\n      type: float\\n  required:\\n  - a\\n  - b\\n\", \"func_name: search\\nfunc_desc: 'search(query: str) -> List[str]\\n\\n  Search for query and return a list of results.'\\nfunc_parameters:\\n  type: object\\n  properties:\\n    query:\\n      type: str\\n  required:\\n  - query\\n\", \"func_name: numpy_sum\\nfunc_desc: 'numpy_sum(arr: numpy.ndarray) -> float\\n\\n  Sum the elements of an array.'\\nfunc_parameters:\\n  type: object\\n  properties:\\n    arr:\\n      type: ndarray\\n  required:\\n  - arr\\n\", \"func_name: add_points\\nfunc_desc: 'add_points(p1: __main__.Point, p2: __main__.Point) -> __main__.Point\\n\\n  None'\\nfunc_parameters:\\n  type: object\\n  properties:\\n    p1:\\n      type: Point\\n      properties:\\n        x:\\n          type: int\\n        y:\\n          type: int\\n      required:\\n      - x\\n      - y\\n    p2:\\n      type: Point\\n      properties:\\n        x:\\n          type: int\\n        y:\\n          type: int\\n      required:\\n      - x\\n      - y\\n  required:\\n  - p1\\n  - p2\\n\"], 'output_format_str': 'Your output should be formatted as a standard JSON instance with the following schema:\\n```\\n{\\n    \"name\": \"The name of the function (str) (optional)\",\\n    \"kwargs\": \"The keyword arguments of the function (Optional) (optional)\"\\n}\\n```\\n-Make sure to always enclose the JSON output in triple backticks (```). Please do not add anything other than valid JSON output!\\n-Use double quotes for the keys and string values.\\n-Follow the JSON formatting conventions.'}, prompt_variables: ['output_format_str', 'tools', 'input_str']\n",
       "  )\n",
       "  (model_client): OpenAIClient()\n",
       "  (output_processors): JsonOutputParser(\n",
       "    data_class_for_json=<class 'lightrag.core.types.Function'>\n",
       "    (json_output_format_prompt): Prompt(\n",
       "      template: Your output should be formatted as a standard JSON instance with the following schema:\n",
       "      ```\n",
       "      {{schema}}\n",
       "      ```\n",
       "      {% if example %}\n",
       "      Here is an example:\n",
       "      ```\n",
       "      {{example}}\n",
       "      ```\n",
       "      {% endif %}\n",
       "      -Make sure to always enclose the JSON output in triple backticks (```). Please do not add anything other than valid JSON output!\n",
       "      -Use double quotes for the keys and string values.\n",
       "      -Follow the JSON formatting conventions., prompt_variables: ['example', 'schema']\n",
       "    )\n",
       "    (output_processors): JsonParser()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create the generator\n",
    "from adalflow.core.generator import Generator\n",
    "from adalflow.core.types import ModelClientType\n",
    "\n",
    "model_kwargs = {\"model\": \"gpt-3.5-turbo\"}\n",
    "prompt_kwargs = {\n",
    "    \"tools\": tool_manager.yaml_definitions,\n",
    "    \"output_format_str\": func_parser.format_instructions(exclude=[\"thought\", \"args\"]),\n",
    "}\n",
    "generator = Generator(\n",
    "    model_client=ModelClientType.OPENAI(),\n",
    "    model_kwargs=model_kwargs,\n",
    "    template=template,\n",
    "    prompt_kwargs=prompt_kwargs,\n",
    "    output_processors=func_parser,\n",
    ")\n",
    "generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr = np.array([[1, 2], [3, 4]])\n",
    "numpy_sum(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# call queries\n",
    "queries = [\n",
    "    \"add 2 and 3\",\n",
    "    \"search for something\",\n",
    "    \"add points (1, 2) and (3, 4)\",\n",
    "    \"sum numpy array with arr = np.array([[1, 2], [3, 4]])\",\n",
    "    \"multiply 2 with local variable x\",\n",
    "    \"divide 2 by 3\",\n",
    "    \"Add 5 to variable y\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "0 Query: add 2 and 3\n",
      "--------------------------------------------------\n",
      "Function: Function(thought=None, name='add', args=[], kwargs={'a': 2, 'b': 3})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Function output: <coroutine object to_thread at 0x127cdeb40>'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1 Query: search for something\n",
      "--------------------------------------------------\n",
      "Function: Function(thought=None, name='search', args=[], kwargs={'query': 'something'})\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/h8/nhgbdr4d18x2r49j4pk5z6gw0000gn/T/ipykernel_1588/3588443470.py:10: RuntimeWarning: coroutine 'to_thread' was never awaited\n",
      "  func_output= tool_manager.execute_func(func)\n",
      "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"Function output: <Task pending name='Task-9' coro=<FunctionTool.acall() running at /Users/liyin/Documents/test/LightRAG/lightrag/core/func_tool.py:124>>\""
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "2 Query: add points (1, 2) and (3, 4)\n",
      "--------------------------------------------------\n",
      "Function: Function(thought=None, name='add_points', args=[], kwargs={'p1': {'x': 1, 'y': 2}, 'p2': {'x': 3, 'y': 4}})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Function output: <coroutine object to_thread at 0x127cdf940>'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "3 Query: sum numpy array with arr = np.array([[1, 2], [3, 4]])\n",
      "--------------------------------------------------\n",
      "Function: Function(thought=None, name='numpy_sum', args=[], kwargs={'arr': [[1, 2], [3, 4]]})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Function output: <coroutine object to_thread at 0x127cdfa40>'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "4 Query: multiply 2 with local variable x\n",
      "--------------------------------------------------\n",
      "Function: Function(thought=None, name='multiply', args=[], kwargs={'a': 2, 'b': 'x'})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Function output: <coroutine object to_thread at 0x127cdf940>'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "5 Query: divide 2 by 3\n",
      "--------------------------------------------------\n",
      "Function: Function(thought=None, name='divide', args=[], kwargs={'a': 2.0, 'b': 3.0})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"Function output: <Task pending name='Task-10' coro=<FunctionTool.acall() running at /Users/liyin/Documents/test/LightRAG/lightrag/core/func_tool.py:124>>\""
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "6 Query: Add 5 to variable y\n",
      "--------------------------------------------------\n",
      "Function: Function(thought=None, name='add', args=[], kwargs={'a': 5, 'b': 10})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Function output: <coroutine object to_thread at 0x127cdfb40>'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for idx, query in enumerate(queries):\n",
    "    prompt_kwargs = {\"input_str\": query}\n",
    "    print(f\"\\n{idx} Query: {query}\")\n",
    "    print(f\"{'-'*50}\")\n",
    "    try:\n",
    "        result = generator(prompt_kwargs=prompt_kwargs)\n",
    "        # print(f\"LLM raw output: {result.raw_response}\")\n",
    "        func = Function.from_dict(result.data)\n",
    "        print(f\"Function: {func}\")\n",
    "        func_output = tool_manager.execute_func(func)\n",
    "        display(f\"Function output: {func_output}\")\n",
    "    except Exception as e:\n",
    "        print(\n",
    "            f\"Failed to execute the function for query: {query}, func: {result.data}, error: {e}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Problems with Function directly:\n",
    "1. difficult to support data types. Unless to update the function to use dict version of the data types to do it.\n",
    "\n",
    "```python\n",
    "def add_points(p1: dict, p2: dict) -> dict:\n",
    "    p1 = Point(**p1)\n",
    "    p2 = Point(**p2)\n",
    "    return add_points_tool.fn(p1, p2).__dict__\n",
    "```\n",
    "2. difficult to use variable as arguments. [TODO: find a proper way to demonstrate it]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "action: Formatted as FuncName(<args>, <kwargs>), where FuncName is the function name, <args> are positional arguments, and <kwargs> are keyword arguments in key=value form. Example: 'FuncName(1, b=2)' calls 'FuncName' with positional argument 1 and keyword argument b=2. (str) (required)\n",
      "FunctionExpression(thought=None, action='add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))')\n"
     ]
    }
   ],
   "source": [
    "# let's use FunctionExpression to call the function instead\n",
    "\n",
    "from adalflow.core.types import FunctionExpression\n",
    "\n",
    "output_data_class = FunctionExpression\n",
    "output_format_str = output_data_class.to_yaml_signature(exclude=[\"thought\"])\n",
    "print(output_format_str)\n",
    "\n",
    "# lets' add one example to be more robust that they should call it with function call expression\n",
    "example = FunctionExpression.from_function(\n",
    "    thought=None, func=add_points, **{\"p1\": Point(1, 2), \"p2\": Point(3, 4)}\n",
    ")\n",
    "print(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt:\n",
      "\n",
      "<SYS>You have these tools available:\n",
      "<TOOLS>\n",
      "1.\n",
      "func_name: multiply\n",
      "func_desc: Multiply two numbers.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: int\n",
      "    b:\n",
      "      type: int\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "2.\n",
      "func_name: add\n",
      "func_desc: Add two numbers.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: int\n",
      "    b:\n",
      "      type: int\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "3.\n",
      "func_name: divide\n",
      "func_desc: Divide two numbers.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: float\n",
      "    b:\n",
      "      type: float\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "4.\n",
      "func_name: search\n",
      "func_desc: Search for query and return a list of results.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    query:\n",
      "      type: str\n",
      "  required:\n",
      "  - query\n",
      "\n",
      "------------------------\n",
      "5.\n",
      "func_name: numpy_sum\n",
      "func_desc: Sum the elements of an array.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    arr:\n",
      "      type: ndarray\n",
      "  required:\n",
      "  - arr\n",
      "\n",
      "------------------------\n",
      "6.\n",
      "func_name: add_points\n",
      "func_desc: None\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    p1:\n",
      "      type: Point\n",
      "      properties:\n",
      "        x:\n",
      "          type: int\n",
      "        y:\n",
      "          type: int\n",
      "      required:\n",
      "      - x\n",
      "      - y\n",
      "    p2:\n",
      "      type: Point\n",
      "      properties:\n",
      "        x:\n",
      "          type: int\n",
      "        y:\n",
      "          type: int\n",
      "      required:\n",
      "      - x\n",
      "      - y\n",
      "  required:\n",
      "  - p1\n",
      "  - p2\n",
      "\n",
      "------------------------\n",
      "</TOOLS>\n",
      "<OUTPUT_FORMAT>\n",
      "Your output should be formatted as a standard JSON instance with the following schema:\n",
      "```\n",
      "{\n",
      "    \"action\": \"Formatted as FuncName(<args>, <kwargs>), where FuncName is the function name, <args> are positional arguments, and <kwargs> are keyword arguments in key=value form. Example: 'FuncName(1, b=2)' calls 'FuncName' with positional argument 1 and keyword argument b=2. (str) (required)\"\n",
      "}\n",
      "```\n",
      "Here is an example:\n",
      "```\n",
      "{\n",
      "    \"action\": \"add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))\"\n",
      "}\n",
      "```\n",
      "-Make sure to always enclose the JSON output in triple backticks (```). Please do not add anything other than valid JSON output!\n",
      "-Use double quotes for the keys and string values.\n",
      "-Follow the JSON formatting conventions.\n",
      "</OUTPUT_FORMAT>\n",
      "</SYS>\n",
      "User: None\n",
      "You:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# also use json output parser and create a new generator\n",
    "\n",
    "parser = JsonOutputParser(data_class=FunctionExpression, example=example)\n",
    "instructions = parser.format_instructions(exclude=[\"thought\"])\n",
    "\n",
    "prompt_kwargs = {\n",
    "    \"tools\": [tool.definition.to_yaml() for tool in tools],\n",
    "    \"output_format_str\": parser.format_instructions(exclude=[\"thought\"]),\n",
    "}\n",
    "generator = Generator(\n",
    "    model_client=ModelClientType.OPENAI(),\n",
    "    model_kwargs=model_kwargs,\n",
    "    template=template,\n",
    "    prompt_kwargs=prompt_kwargs,\n",
    "    output_processors=parser,\n",
    ")\n",
    "\n",
    "generator.print_prompt(**prompt_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Errpr at sandbox_exec: invalid syntax (<string>, line 2)\n",
      "Sandbox error: invalid syntax (<string>, line 2)\n"
     ]
    }
   ],
   "source": [
    "import threading\n",
    "\n",
    "# Define a list of safe built-ins\n",
    "SAFE_BUILTINS = {\n",
    "    \"abs\": abs,\n",
    "    \"all\": all,\n",
    "    \"any\": any,\n",
    "    \"bin\": bin,\n",
    "    \"bool\": bool,\n",
    "    \"bytearray\": bytearray,\n",
    "    \"bytes\": bytes,\n",
    "    \"callable\": callable,\n",
    "    \"chr\": chr,\n",
    "    \"complex\": complex,\n",
    "    \"dict\": dict,\n",
    "    \"divmod\": divmod,\n",
    "    \"enumerate\": enumerate,\n",
    "    \"filter\": filter,\n",
    "    \"float\": float,\n",
    "    \"format\": format,\n",
    "    \"frozenset\": frozenset,\n",
    "    \"getattr\": getattr,\n",
    "    \"hasattr\": hasattr,\n",
    "    \"hash\": hash,\n",
    "    \"hex\": hex,\n",
    "    \"int\": int,\n",
    "    \"isinstance\": isinstance,\n",
    "    \"issubclass\": issubclass,\n",
    "    \"iter\": iter,\n",
    "    \"len\": len,\n",
    "    \"list\": list,\n",
    "    \"map\": map,\n",
    "    \"max\": max,\n",
    "    \"min\": min,\n",
    "    \"next\": next,\n",
    "    \"object\": object,\n",
    "    \"oct\": oct,\n",
    "    \"ord\": ord,\n",
    "    \"pow\": pow,\n",
    "    \"range\": range,\n",
    "    \"repr\": repr,\n",
    "    \"reversed\": reversed,\n",
    "    \"round\": round,\n",
    "    \"set\": set,\n",
    "    \"slice\": slice,\n",
    "    \"sorted\": sorted,\n",
    "    \"str\": str,\n",
    "    \"sum\": sum,\n",
    "    \"tuple\": tuple,\n",
    "    \"type\": type,\n",
    "    \"zip\": zip,\n",
    "}\n",
    "\n",
    "\n",
    "# Define a context manager to limit execution time\n",
    "# Create a sandbox execution function\n",
    "def sandbox_exec(code, context=SAFE_BUILTINS, timeout=5):\n",
    "\n",
    "    try:\n",
    "        compiled_code = compile(code, \"<string>\", \"exec\")\n",
    "\n",
    "        # Result dictionary to store execution results\n",
    "        result = {\"output\": None, \"error\": None}\n",
    "\n",
    "        # Define a target function for the thread\n",
    "        def target():\n",
    "            try:\n",
    "                # Execute the code\n",
    "                exec(compiled_code, context, result)\n",
    "            except Exception as e:\n",
    "                result[\"error\"] = e\n",
    "\n",
    "        # Create a thread to execute the code\n",
    "        thread = threading.Thread(target=target)\n",
    "        thread.start()\n",
    "        thread.join(timeout)\n",
    "\n",
    "        # Check if the thread is still alive (timed out)\n",
    "        if thread.is_alive():\n",
    "            result[\"error\"] = TimeoutError(\"Execution timed out\")\n",
    "            raise TimeoutError(\"Execution timed out\")\n",
    "    except Exception as e:\n",
    "        print(f\"Errpr at sandbox_exec: {e}\")\n",
    "        raise e\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "# Example usage\n",
    "code = \"\"\"\n",
    "def add(a, b+5):\n",
    "    return a + b\n",
    "\n",
    "output = add(1, 2+y)\n",
    "\"\"\"\n",
    "\n",
    "try:\n",
    "    result = sandbox_exec(code)\n",
    "    print(\"Sandbox output:\", result)\n",
    "except TimeoutError as e:\n",
    "    print(e)\n",
    "except Exception as e:\n",
    "    print(\"Sandbox error:\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'multiply': <function multiply at 0x110990900>, 'add': <function add at 0x1108705e0>, 'divide': <function divide at 0x11081ff60>, 'search': <function search at 0x11081fec0>, 'numpy_sum': <function numpy_sum at 0x11190c540>, 'add_points': <function add_points at 0x11190e5c0>, 'Point': <class '__main__.Point'>, 'np': <module 'numpy' from '/Users/liyin/Documents/test/LightRAG/.venv/lib/python3.11/site-packages/numpy/__init__.py'>, 'np.ndarray': <class 'numpy.ndarray'>, 'array': <built-in function array>, 'arr': array([[1, 2],\n",
      "       [3, 4]]), 'np.array': <built-in function array>, 'x': 2}\n",
      "Query: add 2 and 3\n",
      "GeneratorOutput(data={'action': 'add(2, b=3)'}, error=None, usage=None, raw_response='{\\n    \"action\": \"add(2, b=3)\"\\n}')\n",
      "FunctionExpression(thought=None, action='add(2, b=3)')\n",
      "Function(thought=None, name='add', args=[2], kwargs={'b': 3})\n",
      "func output: 5\n",
      "func expr: add(2, b=3)\n",
      "func output: 5\n",
      "sandbox output: {'output': 5, 'error': None}\n",
      "Query: search for something\n",
      "GeneratorOutput(data={'action': \"search('something')\"}, error=None, usage=None, raw_response='{\\n    \"action\": \"search(\\'something\\')\"\\n}')\n",
      "FunctionExpression(thought=None, action=\"search('something')\")\n",
      "Function(thought=None, name='search', args=['something'], kwargs={})\n",
      "func output: <coroutine object search at 0x1225db920>\n",
      "func expr: search('something')\n",
      "func output: <coroutine object search at 0x14c28ea80>\n",
      "sandbox output: {'output': <coroutine object search at 0x1225db920>, 'error': None}\n",
      "Query: add points (1, 2) and (3, 4)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/h8/nhgbdr4d18x2r49j4pk5z6gw0000gn/T/ipykernel_30203/1940296491.py:40: RuntimeWarning: coroutine 'search' was never awaited\n",
      "  fun_output = eval(func_expr.action)\n",
      "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/var/folders/h8/nhgbdr4d18x2r49j4pk5z6gw0000gn/T/ipykernel_30203/1940296491.py:22: RuntimeWarning: coroutine 'search' was never awaited\n",
      "  result = generator(prompt_kwargs=prompt_kwargs)\n",
      "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/var/folders/h8/nhgbdr4d18x2r49j4pk5z6gw0000gn/T/ipykernel_30203/1940296491.py:33: RuntimeWarning: coroutine 'search' was never awaited\n",
      "  fun_output = all_functions_dict[func.name](*func.args, **func.kwargs)\n",
      "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GeneratorOutput(data={'action': 'add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))'}, error=None, usage=None, raw_response='```\\n{\\n    \"action\": \"add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))\"\\n}\\n```')\n",
      "FunctionExpression(thought=None, action='add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))')\n",
      "Function(thought=None, name='add_points', args=[], kwargs={'p1': Point(x=1, y=2), 'p2': Point(x=3, y=4)})\n",
      "func output: Point(x=4, y=6)\n",
      "func expr: add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))\n",
      "func output: Point(x=4, y=6)\n",
      "sandbox output: {'output': Point(x=4, y=6), 'error': None}\n",
      "Query: sum numpy array with arr = np.array([[1, 2], [3, 4]])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Error Field elements must be 2- or 3-tuples, got '3' parsing function call expression: numpy_sum(arr=np.array([1, 2], [3, 4]))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GeneratorOutput(data={'action': 'numpy_sum(arr=np.array([1, 2], [3, 4]))'}, error=None, usage=None, raw_response='{\\n    \"action\": \"numpy_sum(arr=np.array([[1, 2], [3, 4]]))\"\\n}')\n",
      "FunctionExpression(thought=None, action='numpy_sum(arr=np.array([1, 2], [3, 4]))')\n",
      "Error Field elements must be 2- or 3-tuples, got '3' parsing function call expression: numpy_sum(arr=np.array([1, 2], [3, 4]))\n",
      "Failed to execute the function for query: sum numpy array with arr = np.array([[1, 2], [3, 4]]), func: {'action': 'numpy_sum(arr=np.array([1, 2], [3, 4]))'}, error: Error Field elements must be 2- or 3-tuples, got '3' parsing function call expression: numpy_sum(arr=np.array([1, 2], [3, 4]))\n",
      "Field elements must be 2- or 3-tuples, got '3'\n",
      "Failed to execute the function for query: sum numpy array with arr = np.array([[1, 2], [3, 4]]), func: {'action': 'numpy_sum(arr=np.array([1, 2], [3, 4]))'}, error: Field elements must be 2- or 3-tuples, got '3'\n",
      "Query: multiply 2 with local variable x\n",
      "GeneratorOutput(data={'action': 'multiply(2, b=x)'}, error=None, usage=None, raw_response='{\\n    \"action\": \"multiply(2, b=x)\"\\n}')\n",
      "FunctionExpression(thought=None, action='multiply(2, b=x)')\n",
      "Function(thought=None, name='multiply', args=[2], kwargs={'b': 2})\n",
      "func output: 4\n",
      "func expr: multiply(2, b=x)\n",
      "func output: 4\n",
      "sandbox output: {'output': 4, 'error': None}\n",
      "Query: divide 2 by 3\n",
      "GeneratorOutput(data={'action': 'divide(2.0, b=3.0)'}, error=None, usage=None, raw_response='{\\n    \"action\": \"divide(2.0, b=3.0)\"\\n}')\n",
      "FunctionExpression(thought=None, action='divide(2.0, b=3.0)')\n",
      "Function(thought=None, name='divide', args=[2.0], kwargs={'b': 3.0})\n",
      "func output: 0.6666666666666666\n",
      "func expr: divide(2.0, b=3.0)\n",
      "func output: 0.6666666666666666\n",
      "sandbox output: {'output': 0.6666666666666666, 'error': None}\n",
      "Query: Add 5 to variable y\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Error Error: 'y', y does not exist in the context_map. parsing function call expression: add(a=5, b=y)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GeneratorOutput(data={'action': 'add(a=5, b=y)'}, error=None, usage=None, raw_response='{\\n    \"action\": \"add(a=5, b=y)\"\\n}')\n",
      "FunctionExpression(thought=None, action='add(a=5, b=y)')\n",
      "Error Error: 'y', y does not exist in the context_map. parsing function call expression: add(a=5, b=y)\n",
      "Failed to execute the function for query: Add 5 to variable y, func: {'action': 'add(a=5, b=y)'}, error: Error Error: 'y', y does not exist in the context_map. parsing function call expression: add(a=5, b=y)\n",
      "func output: 9\n",
      "sandbox output: {'output': None, 'error': NameError(\"name 'y' is not defined\")}\n"
     ]
    }
   ],
   "source": [
    "# run the generator but we will use FunctionTool.parse_function_call_expr and have a context map\n",
    "\n",
    "all_functions_dict.update(\n",
    "    {\n",
    "        \"Point\": Point,\n",
    "        # support numpy\n",
    "        \"np\": np,\n",
    "        \"np.ndarray\": np.ndarray,\n",
    "        \"array\": np.array,\n",
    "        \"arr\": arr,\n",
    "        \"np.array\": np.array,\n",
    "        \"x\": x,\n",
    "    }\n",
    ")\n",
    "y = 4\n",
    "print(all_functions_dict)\n",
    "for query in queries + [\"Add 5 to variable y\"]:\n",
    "\n",
    "    try:\n",
    "        print(f\"Query: {query}\")\n",
    "        prompt_kwargs = {\"input_str\": query}\n",
    "        result = generator(prompt_kwargs=prompt_kwargs)\n",
    "        print(result)\n",
    "\n",
    "        func_expr = FunctionExpression.from_dict(result.data)\n",
    "\n",
    "        print(func_expr)\n",
    "        assert isinstance(\n",
    "            func_expr, FunctionExpression\n",
    "        ), f\"Expected FunctionExpression, got {type(result.data)}\"\n",
    "\n",
    "        # more secure way to handle function call\n",
    "        func: Function = FunctionTool.parse_function_call_expr(\n",
    "            expr=func_expr.action, context_map=all_functions_dict\n",
    "        )\n",
    "        print(func)\n",
    "        fun_output = all_functions_dict[func.name](*func.args, **func.kwargs)\n",
    "        print(\"func output:\", fun_output)\n",
    "\n",
    "        print(f\"func expr: {func_expr.action}\")\n",
    "\n",
    "        # eval without security check by using eval directly\n",
    "        # less secure but even more powerful and flexible\n",
    "        fun_output = eval(func_expr.action)\n",
    "        print(\"func output:\", fun_output)\n",
    "\n",
    "        # sandbox_exec\n",
    "        action = \"output=\" + func_expr.action\n",
    "        result = sandbox_exec(action, context={**SAFE_BUILTINS, **all_functions_dict})\n",
    "        print(\"sandbox output:\", result)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\n",
    "            f\"Failed to execute the function for query: {query}, func: {result.data}, error: {e}\"\n",
    "        )\n",
    "        try:\n",
    "            fun_output = eval(func_expr.action)\n",
    "            print(\"func output:\", fun_output)\n",
    "\n",
    "            # sandbox_exec\n",
    "            action = \"output=\" + func_expr.action\n",
    "            result = sandbox_exec(\n",
    "                action, context={**SAFE_BUILTINS, **all_functions_dict}\n",
    "            )\n",
    "            print(\"sandbox output:\", result)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(\n",
    "                f\"Failed to execute the function for query: {query}, func: {result.data}, error: {e}\"\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multiple function calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "multple_function_call_template = r\"\"\"<SYS>You have these tools available:\n",
    "{% if tools %}\n",
    "<TOOLS>\n",
    "{% for tool in tools %}\n",
    "{{ loop.index }}.\n",
    "{{tool}}\n",
    "------------------------\n",
    "{% endfor %}\n",
    "</TOOLS>\n",
    "{% endif %}\n",
    "<OUTPUT_FORMAT>\n",
    "Here is how you call one function.\n",
    "{{output_format_str}}\n",
    "Return a List using `[]` of the above JSON objects. You can have length of 1 or more.\n",
    "Do not call multiple functions in one action field.\n",
    "</OUTPUT_FORMAT>\n",
    "<SYS>\n",
    "{{input_str}}\n",
    "You:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generator(\n",
      "  model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
      "  (prompt): Prompt(\n",
      "    template: <SYS>You have these tools available:\n",
      "    {% if tools %}\n",
      "    <TOOLS>\n",
      "    {% for tool in tools %}\n",
      "    {{ loop.index }}.\n",
      "    {{tool}}\n",
      "    ------------------------\n",
      "    {% endfor %}\n",
      "    </TOOLS>\n",
      "    {% endif %}\n",
      "    <OUTPUT_FORMAT>\n",
      "    Here is how you call one function.\n",
      "    {{output_format_str}}\n",
      "    Return a List using `[]` of the above JSON objects. You can have length of 1 or more.\n",
      "    Do not call multiple functions in one action field.\n",
      "    </OUTPUT_FORMAT>\n",
      "    <SYS>\n",
      "    {{input_str}}\n",
      "    You:\n",
      "    , prompt_kwargs: {'tools': ['func_name: multiply\\nfunc_desc: Multiply two numbers.\\nfunc_parameters:\\n  type: object\\n  properties:\\n    a:\\n      type: int\\n    b:\\n      type: int\\n  required:\\n  - a\\n  - b\\n', 'func_name: add\\nfunc_desc: Add two numbers.\\nfunc_parameters:\\n  type: object\\n  properties:\\n    a:\\n      type: int\\n    b:\\n      type: int\\n  required:\\n  - a\\n  - b\\n', 'func_name: divide\\nfunc_desc: Divide two numbers.\\nfunc_parameters:\\n  type: object\\n  properties:\\n    a:\\n      type: float\\n    b:\\n      type: float\\n  required:\\n  - a\\n  - b\\n', 'func_name: search\\nfunc_desc: Search for query and return a list of results.\\nfunc_parameters:\\n  type: object\\n  properties:\\n    query:\\n      type: str\\n  required:\\n  - query\\n', 'func_name: numpy_sum\\nfunc_desc: Sum the elements of an array.\\nfunc_parameters:\\n  type: object\\n  properties:\\n    arr:\\n      type: ndarray\\n  required:\\n  - arr\\n', 'func_name: add_points\\nfunc_desc: None\\nfunc_parameters:\\n  type: object\\n  properties:\\n    p1:\\n      type: Point\\n      properties:\\n        x:\\n          type: int\\n        y:\\n          type: int\\n      required:\\n      - x\\n      - y\\n    p2:\\n      type: Point\\n      properties:\\n        x:\\n          type: int\\n        y:\\n          type: int\\n      required:\\n      - x\\n      - y\\n  required:\\n  - p1\\n  - p2\\n'], 'output_format_str': 'Your output should be formatted as a standard JSON instance with the following schema:\\n```\\n{\\n    \"action\": \"Formatted as FuncName(<args>, <kwargs>), where FuncName is the function name, <args> are positional arguments, and <kwargs> are keyword arguments in key=value form. Example: \\'FuncName(1, b=2)\\' calls \\'FuncName\\' with positional argument 1 and keyword argument b=2. (str) (required)\"\\n}\\n```\\nHere is an example:\\n```\\n{\\n    \"action\": \"add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))\"\\n}\\n```\\n-Make sure to always enclose the JSON output in triple backticks (```). Please do not add anything other than valid JSON output!\\n-Use double quotes for the keys and string values.\\n-Follow the JSON formatting conventions.'}, prompt_variables: ['output_format_str', 'input_str', 'tools']\n",
      "  )\n",
      "  (model_client): OpenAIClient()\n",
      "  (output_processors): JsonParser()\n",
      ")\n",
      "Prompt:\n",
      "\n",
      "<SYS>You have these tools available:\n",
      "<TOOLS>\n",
      "1.\n",
      "func_name: multiply\n",
      "func_desc: Multiply two numbers.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: int\n",
      "    b:\n",
      "      type: int\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "2.\n",
      "func_name: add\n",
      "func_desc: Add two numbers.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: int\n",
      "    b:\n",
      "      type: int\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "3.\n",
      "func_name: divide\n",
      "func_desc: Divide two numbers.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    a:\n",
      "      type: float\n",
      "    b:\n",
      "      type: float\n",
      "  required:\n",
      "  - a\n",
      "  - b\n",
      "\n",
      "------------------------\n",
      "4.\n",
      "func_name: search\n",
      "func_desc: Search for query and return a list of results.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    query:\n",
      "      type: str\n",
      "  required:\n",
      "  - query\n",
      "\n",
      "------------------------\n",
      "5.\n",
      "func_name: numpy_sum\n",
      "func_desc: Sum the elements of an array.\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    arr:\n",
      "      type: ndarray\n",
      "  required:\n",
      "  - arr\n",
      "\n",
      "------------------------\n",
      "6.\n",
      "func_name: add_points\n",
      "func_desc: None\n",
      "func_parameters:\n",
      "  type: object\n",
      "  properties:\n",
      "    p1:\n",
      "      type: Point\n",
      "      properties:\n",
      "        x:\n",
      "          type: int\n",
      "        y:\n",
      "          type: int\n",
      "      required:\n",
      "      - x\n",
      "      - y\n",
      "    p2:\n",
      "      type: Point\n",
      "      properties:\n",
      "        x:\n",
      "          type: int\n",
      "        y:\n",
      "          type: int\n",
      "      required:\n",
      "      - x\n",
      "      - y\n",
      "  required:\n",
      "  - p1\n",
      "  - p2\n",
      "\n",
      "------------------------\n",
      "</TOOLS>\n",
      "<OUTPUT_FORMAT>\n",
      "Here is how you call one function.\n",
      "Your output should be formatted as a standard JSON instance with the following schema:\n",
      "```\n",
      "{\n",
      "    \"action\": \"Formatted as FuncName(<args>, <kwargs>), where FuncName is the function name, <args> are positional arguments, and <kwargs> are keyword arguments in key=value form. Example: 'FuncName(1, b=2)' calls 'FuncName' with positional argument 1 and keyword argument b=2. (str) (required)\"\n",
      "}\n",
      "```\n",
      "Here is an example:\n",
      "```\n",
      "{\n",
      "    \"action\": \"add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))\"\n",
      "}\n",
      "```\n",
      "-Make sure to always enclose the JSON output in triple backticks (```). Please do not add anything other than valid JSON output!\n",
      "-Use double quotes for the keys and string values.\n",
      "-Follow the JSON formatting conventions.\n",
      "Return a List using `[]` of the above JSON objects. You can have length of 1 or more.\n",
      "Do not call multiple functions in one action field.\n",
      "</OUTPUT_FORMAT>\n",
      "<SYS>\n",
      "None\n",
      "You:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "queries = [\n",
    "    \"add 2 and 3\",\n",
    "    \"search for something\",\n",
    "    \"add points (1, 2) and (3, 4)\",\n",
    "    \"sum numpy array with arr = np.array([[1, 2], [3, 4]])\",\n",
    "    \"multiply 2 with local variable x\",\n",
    "    \"divide 2 by 3\",\n",
    "]\n",
    "\n",
    "from adalflow.core.string_parser import JsonParser  # improve a list of json\n",
    "\n",
    "preset_prompt_kwargs = {\n",
    "    \"tools\": [tool.definition.to_yaml() for tool in tools],\n",
    "    \"output_format_str\": parser.format_instructions(exclude=[\"thought\"]),\n",
    "}\n",
    "multi_call_gen = Generator(\n",
    "    model_client=ModelClientType.OPENAI(),\n",
    "    model_kwargs=model_kwargs,\n",
    "    template=multple_function_call_template,\n",
    "    prompt_kwargs=preset_prompt_kwargs,\n",
    "    output_processors=JsonParser(),\n",
    ")\n",
    "print(multi_call_gen)\n",
    "multi_call_gen.print_prompt()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query: add 2 and 3 and search for something\n",
      "_________________________\n",
      "\n",
      "GeneratorOutput(data={'action': 'add(2, b=3)'}, error=None, usage=None, raw_response='```  \\n{\\n    \"action\": \"add(2, b=3)\"\\n}\\n{\\n    \"action\": \"search(query=\\'something\\')\"\\n}\\n```')\n",
      "'str' object has no attribute 'items'\n",
      "Failed to parse the function for query: add 2 and 3 and search for something, func: {'action': 'add(2, b=3)'}, error: 'str' object has no attribute 'items'\n",
      "Query: add points (1, 2) and (3, 4) and sum numpy array with arr = np.array([[1, 2], [3, 4]])\n",
      "_________________________\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Error Field elements must be 2- or 3-tuples, got '3' parsing function call expression: numpy_sum(arr=np.array([1, 2], [3, 4]))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GeneratorOutput(data=[{'action': 'add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))'}, {'action': 'numpy_sum(arr=np.array([1, 2], [3, 4]))'}], error=None, usage=None, raw_response='```json\\n[\\n    {\\n        \"action\": \"add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))\"\\n    },\\n    {\\n        \"action\": \"numpy_sum(arr=np.array([[1, 2], [3, 4]]))\"\\n    }\\n]\\n```')\n",
      "[FunctionExpression(thought=None, action='add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))'), FunctionExpression(thought=None, action='numpy_sum(arr=np.array([1, 2], [3, 4]))')]\n",
      "Function(thought=None, name='add_points', args=[], kwargs={'p1': Point(x=1, y=2), 'p2': Point(x=3, y=4)})\n",
      "func output: Point(x=4, y=6)\n",
      "Error Field elements must be 2- or 3-tuples, got '3' parsing function call expression: numpy_sum(arr=np.array([1, 2], [3, 4]))\n",
      "Failed to execute the function for query: add points (1, 2) and (3, 4) and sum numpy array with arr = np.array([[1, 2], [3, 4]]), func: [{'action': 'add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))'}, {'action': 'numpy_sum(arr=np.array([1, 2], [3, 4]))'}], error: Error Field elements must be 2- or 3-tuples, got '3' parsing function call expression: numpy_sum(arr=np.array([1, 2], [3, 4]))\n",
      "func expr: add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))\n",
      "func output: Point(x=4, y=6)\n",
      "func expr: numpy_sum(arr=np.array([1, 2], [3, 4]))\n",
      "Field elements must be 2- or 3-tuples, got '3'\n",
      "Failed to execute the function for query: add points (1, 2) and (3, 4) and sum numpy array with arr = np.array([[1, 2], [3, 4]]), func: [{'action': 'add_points(p1=Point(x=1, y=2), p2=Point(x=3, y=4))'}, {'action': 'numpy_sum(arr=np.array([1, 2], [3, 4]))'}], error: Field elements must be 2- or 3-tuples, got '3'\n",
      "sandbox output: {'output': Point(x=4, y=6), 'error': None}\n",
      "sandbox output: {'output': None, 'error': TypeError(\"Field elements must be 2- or 3-tuples, got '3'\")}\n",
      "fun_output by sandbox: [{'output': Point(x=4, y=6), 'error': None}, {'output': None, 'error': TypeError(\"Field elements must be 2- or 3-tuples, got '3'\")}]\n",
      "_________________________\n",
      "\n",
      "Query: multiply 2 with local variable x and divide 2 by 3\n",
      "_________________________\n",
      "\n",
      "GeneratorOutput(data=[{'action': 'multiply(2, b=x)'}, {'action': 'divide(a=2, b=3)'}], error=None, usage=None, raw_response='```json\\n[\\n    {\\n        \"action\": \"multiply(2, b=x)\"\\n    },\\n    {\\n        \"action\": \"divide(a=2, b=3)\"\\n    }\\n]\\n```')\n",
      "[FunctionExpression(thought=None, action='multiply(2, b=x)'), FunctionExpression(thought=None, action='divide(a=2, b=3)')]\n",
      "Function(thought=None, name='multiply', args=[2], kwargs={'b': 2})\n",
      "func output: 4\n",
      "Function(thought=None, name='divide', args=[], kwargs={'a': 2, 'b': 3})\n",
      "func output: 0.6666666666666666\n",
      "fun_output by parsing: [4, 0.6666666666666666]\n",
      "_________________________\n",
      "\n",
      "func expr: multiply(2, b=x)\n",
      "func output: 4\n",
      "func expr: divide(a=2, b=3)\n",
      "func output: 0.6666666666666666\n",
      "fun_output by eval: [4, 0.6666666666666666]\n",
      "_________________________\n",
      "\n",
      "sandbox output: {'output': 4, 'error': None}\n",
      "sandbox output: {'output': 0.6666666666666666, 'error': None}\n",
      "fun_output by sandbox: [{'output': 4, 'error': None}, {'output': 0.6666666666666666, 'error': None}]\n",
      "_________________________\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/h8/nhgbdr4d18x2r49j4pk5z6gw0000gn/T/ipykernel_30203/2050537243.py:47: RuntimeWarning: coroutine 'search' was never awaited\n",
      "  func_outputs_1 = [execute_function_by_parsing(func_expr, all_functions_dict) for func_expr in func_exprs]\n",
      "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/var/folders/h8/nhgbdr4d18x2r49j4pk5z6gw0000gn/T/ipykernel_30203/2050537243.py:55: RuntimeWarning: coroutine 'search' was never awaited\n",
      "  func_outputs_2 = [execute_function_by_eval(func_expr) for func_expr in func_exprs]\n",
      "RuntimeWarning: Enable tracemalloc to get the object allocation traceback\n"
     ]
    }
   ],
   "source": [
    "def execute_function_by_parsing(\n",
    "    func_expr: FunctionExpression, all_functions_dict: Dict[str, Any]\n",
    ") -> Any:\n",
    "    func: Function = FunctionTool.parse_function_call_expr(\n",
    "        expr=func_expr.action, context_map=all_functions_dict\n",
    "    )\n",
    "    print(func)\n",
    "    fun_output = all_functions_dict[func.name](*func.args, **func.kwargs)\n",
    "    print(\"func output:\", fun_output)\n",
    "    return fun_output\n",
    "\n",
    "\n",
    "def execute_function_by_eval(func_expr: FunctionExpression) -> Any:\n",
    "\n",
    "    print(f\"func expr: {func_expr.action}\")\n",
    "\n",
    "    # eval without security check by using eval directly\n",
    "    # less secure but even more powerful and flexible\n",
    "    fun_output = eval(func_expr.action)\n",
    "    print(\"func output:\", fun_output)\n",
    "    return fun_output\n",
    "\n",
    "\n",
    "def execute_function_by_sandbox(\n",
    "    func_expr: FunctionExpression, all_functions_dict: Dict[str, Any]\n",
    ") -> Any:\n",
    "    # sandbox_exec\n",
    "    action = \"output=\" + func_expr.action\n",
    "    result = sandbox_exec(action, context={**SAFE_BUILTINS, **all_functions_dict})\n",
    "    print(\"sandbox output:\", result)\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "for i in range(0, len(queries), 2):\n",
    "    query = \" and \".join(queries[i : i + 2])\n",
    "    print(f\"Query: {query}\\n_________________________\\n\")\n",
    "    prompt_kwargs = {\"input_str\": query}\n",
    "    result = multi_call_gen(prompt_kwargs=prompt_kwargs)\n",
    "    print(result)\n",
    "\n",
    "    try:\n",
    "\n",
    "        func_exprs = [FunctionExpression.from_dict(item) for item in result.data]\n",
    "\n",
    "        print(func_exprs)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\n",
    "            f\"Failed to parse the function for query: {query}, func: {result.data}, error: {e}\"\n",
    "        )\n",
    "        continue\n",
    "    try:\n",
    "        func_outputs_1 = [\n",
    "            execute_function_by_parsing(func_expr, all_functions_dict)\n",
    "            for func_expr in func_exprs\n",
    "        ]\n",
    "        print(f\"fun_output by parsing: {func_outputs_1}\\n_________________________\\n\")\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\n",
    "            f\"Failed to execute the function for query: {query}, func: {result.data}, error: {e}\"\n",
    "        )\n",
    "\n",
    "    try:\n",
    "\n",
    "        func_outputs_2 = [\n",
    "            execute_function_by_eval(func_expr) for func_expr in func_exprs\n",
    "        ]\n",
    "        print(f\"fun_output by eval: {func_outputs_2}\\n_________________________\\n\")\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\n",
    "            f\"Failed to execute the function for query: {query}, func: {result.data}, error: {e}\"\n",
    "        )\n",
    "\n",
    "    try:\n",
    "\n",
    "        func_outputs_3 = [\n",
    "            execute_function_by_sandbox(func_expr, all_functions_dict)\n",
    "            for func_expr in func_exprs\n",
    "        ]\n",
    "        print(f\"fun_output by sandbox: {func_outputs_3}\\n_________________________\\n\")\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\n",
    "            f\"Failed to execute the function for query: {query}, func: {result.data}, error: {e}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "response: ChatCompletion(id='chatcmpl-9eDBpPnQkSDM90VqKgmtGsMJ3k7jJ', choices=[Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, role='assistant', function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_0f6vJCTGXHRDlr3Uns6cczlJ', function=Function(arguments='{\"location\": \"San Francisco, CA\", \"unit\": \"celsius\"}', name='get_current_weather'), type='function'), ChatCompletionMessageToolCall(id='call_GZwXUQ2hVeLmOuf6Ty28JluG', function=Function(arguments='{\"location\": \"Tokyo, Japan\", \"unit\": \"celsius\"}', name='get_current_weather'), type='function'), ChatCompletionMessageToolCall(id='call_ASnPn2pGcyixg7AjCr4DSdYs', function=Function(arguments='{\"location\": \"Paris, France\", \"unit\": \"celsius\"}', name='get_current_weather'), type='function')]))], created=1719370849, model='gpt-4o-2024-05-13', object='chat.completion', service_tier=None, system_fingerprint='fp_4008e3b719', usage=CompletionUsage(completion_tokens=83, prompt_tokens=88, total_tokens=171))\n",
      "tool_calls: [ChatCompletionMessageToolCall(id='call_0f6vJCTGXHRDlr3Uns6cczlJ', function=Function(arguments='{\"location\": \"San Francisco, CA\", \"unit\": \"celsius\"}', name='get_current_weather'), type='function'), ChatCompletionMessageToolCall(id='call_GZwXUQ2hVeLmOuf6Ty28JluG', function=Function(arguments='{\"location\": \"Tokyo, Japan\", \"unit\": \"celsius\"}', name='get_current_weather'), type='function'), ChatCompletionMessageToolCall(id='call_ASnPn2pGcyixg7AjCr4DSdYs', function=Function(arguments='{\"location\": \"Paris, France\", \"unit\": \"celsius\"}', name='get_current_weather'), type='function')]\n",
      "ChatCompletion(id='chatcmpl-9eDBq0oXMtJYAB9TOfAqm9dNNQoJf', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content='Here are the current temperatures in Celsius:\\n\\n- **San Francisco, CA:** \\\\(72^\\\\circ \\\\text{C}\\\\)\\n- **Tokyo, Japan:** \\\\(10^\\\\circ \\\\text{C}\\\\)\\n- **Paris, France:** \\\\(22^\\\\circ \\\\text{C}\\\\)\\n\\nLet me know if you need any more information!', role='assistant', function_call=None, tool_calls=None))], created=1719370850, model='gpt-4o-2024-05-13', object='chat.completion', service_tier=None, system_fingerprint='fp_d576307f90', usage=CompletionUsage(completion_tokens=67, prompt_tokens=176, total_tokens=243))\n"
     ]
    }
   ],
   "source": [
    "# first check the openai's function call apis\n",
    "\n",
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI()\n",
    "\n",
    "\n",
    "# Example dummy function hard coded to return the same weather\n",
    "# In production, this could be your backend API or an external API\n",
    "def get_current_weather(location, unit=\"fahrenheit\"):\n",
    "    \"\"\"Get the current weather in a given location\"\"\"\n",
    "    if \"tokyo\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Tokyo\", \"temperature\": \"10\", \"unit\": unit})\n",
    "    elif \"san francisco\" in location.lower():\n",
    "        return json.dumps(\n",
    "            {\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": unit}\n",
    "        )\n",
    "    elif \"paris\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": unit})\n",
    "    else:\n",
    "        return json.dumps({\"location\": location, \"temperature\": \"unknown\"})\n",
    "\n",
    "\n",
    "def run_conversation():\n",
    "    # Step 1: send the conversation and available functions to the model\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"What's the weather like in San Francisco, Tokyo, and Paris in celsius?\",\n",
    "        }\n",
    "    ]\n",
    "    tools = [\n",
    "        {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": \"get_current_weather\",\n",
    "                \"description\": \"Get the current weather in a given location\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"location\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The city and state, e.g. San Francisco, CA\",\n",
    "                        },\n",
    "                        \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"]},\n",
    "                    },\n",
    "                    \"required\": [\"location\"],\n",
    "                },\n",
    "            },\n",
    "        }\n",
    "    ]\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        messages=messages,\n",
    "        tools=tools,\n",
    "        tool_choice=\"auto\",  # auto is default, but we'll be explicit\n",
    "    )\n",
    "    print(f\"response: {response}\")\n",
    "    response_message = response.choices[0].message\n",
    "    tool_calls = response_message.tool_calls\n",
    "\n",
    "    print(f\"tool_calls: {tool_calls}\")\n",
    "    # Step 2: check if the model wanted to call a function\n",
    "    if tool_calls:\n",
    "        # Step 3: call the function\n",
    "        # Note: the JSON response may not always be valid; be sure to handle errors\n",
    "        available_functions = {\n",
    "            \"get_current_weather\": get_current_weather,\n",
    "        }  # only one function in this example, but you can have multiple\n",
    "        messages.append(response_message)  # extend conversation with assistant's reply\n",
    "        # Step 4: send the info for each function call and function response to the model\n",
    "        for tool_call in tool_calls:\n",
    "            function_name = tool_call.function.name\n",
    "            function_to_call = available_functions[function_name]\n",
    "            function_args = json.loads(\n",
    "                tool_call.function.arguments\n",
    "            )  # use json.loads to convert a string to a dictionary\n",
    "            # function_response = function_to_call(\n",
    "            #     location=function_args.get(\"location\"),\n",
    "            #     unit=function_args.get(\"unit\"),\n",
    "            # )\n",
    "            # you have to exactly know the arguments, this does not make sense. How would i know its arguments. **function_args (makes more sense)\n",
    "            function_response = function_to_call(**function_args)\n",
    "            messages.append(\n",
    "                {\n",
    "                    \"tool_call_id\": tool_call.id,\n",
    "                    \"role\": \"tool\",\n",
    "                    \"name\": function_name,\n",
    "                    \"content\": function_response,\n",
    "                }\n",
    "            )  # extend conversation with function response\n",
    "        second_response = client.chat.completions.create(\n",
    "            model=\"gpt-4o\",\n",
    "            messages=messages,\n",
    "        )  # get a new response from the model where it can see the function response\n",
    "        return second_response\n",
    "\n",
    "\n",
    "print(run_conversation())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function(arguments='{\"location\": \"Tokyo, Japan\", \"unit\": \"celsius\"}', name='get_current_weather'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are two important pieces. Getting function schema is not difficult and can be standarized.\n",
    "\n",
    "The second piece is how to call the function, and how to execute it. The how to call the function depends on how we execute it.\n",
    "\n",
    "How to execute a function:\n",
    "1. Eval (LLM will output the code to call the function (in string format))-> Language generation.\n",
    "2. We manage a function map, and we ask LLm to output either the code string or a structure with the function name and the arguments. We can use the function map to call the function. If its code string, we will have to parse the function call into the name and the arguments. If its a structure, we will have to convert it to data structure that can be used to call the function.\n",
    "\n",
    "There are just so many different ways to do the actual function call, and different LLM might react differetntly in accuracy to each output format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Function(arguments='{\"location\": \"Paris, France\"}', name='get_current_weather'), type='function')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_current_weather(location: str, unit: str = \"fahrenheit\"):\n",
    "    \"\"\"Get the current weather in a given location\"\"\"\n",
    "    if \"tokyo\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Tokyo\", \"temperature\": \"10\", \"unit\": unit})\n",
    "    elif \"san francisco\" in location.lower():\n",
    "        return json.dumps(\n",
    "            {\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": unit}\n",
    "        )\n",
    "    elif \"paris\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": unit})\n",
    "    else:\n",
    "        return json.dumps({\"location\": location, \"temperature\": \"unknown\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v2\n",
    "\n",
    "from adalflow.core.base_data_class import DataClass\n",
    "from dataclasses import dataclass, field\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Weather(DataClass):\n",
    "    location: str = field(\n",
    "        metadata={\"description\": \"The city and state, e.g. San Francisco, CA\"}\n",
    "    )\n",
    "    unit: str = field(metadata={\"enum\": [\"celsius\", \"fahrenheit\"]})\n",
    "\n",
    "\n",
    "def get_current_weather_2(weather: Weather):\n",
    "    \"\"\"Get the current weather in a given location\"\"\"\n",
    "    if \"tokyo\" in weather.location.lower():\n",
    "        return json.dumps(\n",
    "            {\"location\": \"Tokyo\", \"temperature\": \"10\", \"unit\": weather.unit}\n",
    "        )\n",
    "    elif \"san francisco\" in weather.location.lower():\n",
    "        return json.dumps(\n",
    "            {\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": weather.unit}\n",
    "        )\n",
    "    elif \"paris\" in weather.location.lower():\n",
    "        return json.dumps(\n",
    "            {\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": weather.unit}\n",
    "        )\n",
    "    else:\n",
    "        return json.dumps({\"location\": weather.location, \"temperature\": \"unknown\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name: weather, parameter: weather: __main__.Weather    <class '__main__.Weather'>\n",
      "type_hints[name]: <class '__main__.Weather'>\n",
      "name: location, parameter: location: str    <class 'str'>\n",
      "name: unit, parameter: unit: str    <class 'str'>\n",
      "{\n",
      "    \"name\": \"get_current_weather_2\",\n",
      "    \"description\": \"get_current_weather_2(weather: __main__.Weather)\\nGet the current weather in a given location\",\n",
      "    \"parameters\": {\n",
      "        \"type\": \"object\",\n",
      "        \"properties\": {\n",
      "            \"weather\": {\n",
      "                \"type\": \"Weather\",\n",
      "                \"description\": \"The city and state, e.g. San Francisco, CA\",\n",
      "                \"enum\": [\n",
      "                    \"celsius\",\n",
      "                    \"fahrenheit\"\n",
      "                ]\n",
      "            }\n",
      "        },\n",
      "        \"required\": [\n",
      "            \"weather\"\n",
      "        ],\n",
      "        \"definitions\": {\n",
      "            \"weather\": {\n",
      "                \"type\": \"object\",\n",
      "                \"properties\": {\n",
      "                    \"location\": {\n",
      "                        \"type\": \"str\"\n",
      "                    },\n",
      "                    \"unit\": {\n",
      "                        \"type\": \"str\"\n",
      "                    }\n",
      "                },\n",
      "                \"required\": [\n",
      "                    \"location\",\n",
      "                    \"unit\"\n",
      "                ]\n",
      "            }\n",
      "        }\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# Create a tool from the class\n",
    "\n",
    "tool_2 = FunctionTool.from_defaults(fn=get_current_weather_2)\n",
    "\n",
    "print(tool_2.metadata.to_json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Llamaindex\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adalflow_fn_schema = {\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\n",
    "        \"weather\": {\n",
    "            \"type\": \"Weather\",\n",
    "            \"desc\": \"The city and state, e.g. San Francisco, CA\",\n",
    "            \"enum\": [\"celsius\", \"fahrenheit\"],\n",
    "        }\n",
    "    },\n",
    "    \"required\": [\"weather\"],\n",
    "    \"definitions\": {\n",
    "        \"weather\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\"location\": {\"type\": \"str\"}, \"unit\": {\"type\": \"str\"}},\n",
    "            \"required\": [\"location\", \"unit\"],\n",
    "        }\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_fn_schema = {\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\"weather\": {\"$ref\": \"#/definitions/Weather\"}},\n",
    "    \"required\": [\"weather\"],\n",
    "    \"definitions\": {\n",
    "        \"Weather\": {\n",
    "            \"title\": \"Weather\",\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"location\": {\n",
    "                    \"title\": \"Location\",\n",
    "                    \"desc\": \"The city and state, e.g. San Francisco, CA\",\n",
    "                    \"type\": \"string\",\n",
    "                },\n",
    "                \"unit\": {\n",
    "                    \"title\": \"Unit\",\n",
    "                    \"enum\": [\"celsius\", \"fahrenheit\"],\n",
    "                    \"type\": \"string\",\n",
    "                },\n",
    "            },\n",
    "            \"required\": [\"location\", \"unit\"],\n",
    "            \"additionalProperties\": False,\n",
    "        }\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# level 1, call function with default python data types\n",
    "# such as str, int, float, list, dict, etc.\n",
    "def _get_current_weather(location: str, unit: str = \"fahrenheit\"):\n",
    "    \"\"\"Get the current weather in a given location\"\"\"\n",
    "    if \"tokyo\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Tokyo\", \"temperature\": \"10\", \"unit\": unit})\n",
    "    elif \"san francisco\" in location.lower():\n",
    "        return json.dumps(\n",
    "            {\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": unit}\n",
    "        )\n",
    "    elif \"paris\" in location.lower():\n",
    "        return json.dumps({\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": unit})\n",
    "    else:\n",
    "        return json.dumps({\"location\": location, \"temperature\": \"unknown\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FunctionTool(metadata=ToolMetadata(name='_get_current_weather', description=\"_get_current_weather(location: str, unit: str = 'fahrenheit')\\nGet the current weather in a given location\", parameters={'type': 'object', 'properties': {'location': {'type': 'str'}, 'unit': {'type': 'str', 'default': 'fahrenheit'}}, 'required': ['location']}), fn=<function _get_current_weather at 0x10806ff60>, async_fn=None)\n"
     ]
    }
   ],
   "source": [
    "# prepare function tool\n",
    "weather_tool = FunctionTool.from_defaults(fn=_get_current_weather)\n",
    "print(weather_tool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Generator(\n",
       "  model_kwargs={'model': 'gpt-3.5-turbo', 'temperature': 0.3, 'stream': False}, \n",
       "  (prompt): Prompt(\n",
       "    template: <SYS>You have these tools available:\n",
       "        <TOOLS>\n",
       "        {% for tool in tools %}\n",
       "        {{ loop.index }}. ToolName: {{ tool.metadata.name }}\n",
       "            Tool Description: {{ tool.metadata.description }}\n",
       "            Tool Parameters: {{ tool.metadata.fn_schema_str }}   \n",
       "        __________\n",
       "        {% endfor %}\n",
       "        </TOOLS>\n",
       "        {{output_format_str}}\n",
       "        </SYS>\n",
       "        User: {{input_str}}\n",
       "        You:\n",
       "        , prompt_kwargs: {'output_format_str': 'Your output should be formatted as a standard YAML instance with the following schema:\\n```\\nname: The name of the function (str) (required)\\nargs: The arguments of the function (Dict) (required)\\n```\\n\\n-Make sure to always enclose the YAML output in triple backticks (```). Please do not add anything other than valid YAML output!\\n-Follow the YAML formatting conventions with an indent of 2 spaces.\\n-Quote the string values properly.\\n'}, prompt_variables: ['input_str', 'output_format_str', 'tools']\n",
       "  )\n",
       "  (model_client): OpenAIClient()\n",
       "  (output_processors): YamlOutputParser(\n",
       "    data_class_for_yaml=<class '__main__.Function'>\n",
       "    (yaml_output_format_prompt): Prompt(\n",
       "      template: Your output should be formatted as a standard YAML instance with the following schema:\n",
       "      ```\n",
       "      {{schema}}\n",
       "      ```\n",
       "      {% if example %}\n",
       "      Here is an example:\n",
       "      ```\n",
       "      {{example}}\n",
       "      ```\n",
       "      {% endif %}\n",
       "      \n",
       "      -Make sure to always enclose the YAML output in triple backticks (```). Please do not add anything other than valid YAML output!\n",
       "      -Follow the YAML formatting conventions with an indent of 2 spaces.\n",
       "      -Quote the string values properly.\n",
       "      , prompt_variables: ['example', 'schema']\n",
       "    )\n",
       "    (output_processors): YamlParser()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# prepare a minimal function calling template\n",
    "template = r\"\"\"<SYS>You have these tools available:\n",
    "    <TOOLS>\n",
    "    {% for tool in tools %}\n",
    "    {{ loop.index }}. ToolName: {{ tool.metadata.name }}\n",
    "        Tool Description: {{ tool.metadata.description }}\n",
    "        Tool Parameters: {{ tool.metadata.fn_schema_str }}   \n",
    "    __________\n",
    "    {% endfor %}\n",
    "    </TOOLS>\n",
    "    {{output_format_str}}\n",
    "    </SYS>\n",
    "    User: {{input_str}}\n",
    "    You:\n",
    "    \"\"\"\n",
    "\n",
    "multiple_function_call_template = r\"\"\"<SYS>You can answer user query with these tools:\n",
    "    <TOOLS>\n",
    "    {% for tool in tools %}\n",
    "    {{ loop.index }}. ToolName: {{ tool.metadata.name }}\n",
    "        Tool Description: {{ tool.metadata.description }}\n",
    "        Tool Parameters: {{ tool.metadata.fn_schema_str }}   \n",
    "    __________\n",
    "    {% endfor %}\n",
    "    </TOOLS>\n",
    "    You can call multiple tools by return a list of the following format:\n",
    "    {{output_format_str}}\n",
    "    </SYS>\n",
    "    User: {{input_str}}\n",
    "    You:\n",
    "    \"\"\"\n",
    "\n",
    "from typing import Dict, Any\n",
    "from adalflow.core.generator import Generator\n",
    "from adalflow.core.types import ModelClientType\n",
    "from adalflow.components.output_parsers import YamlOutputParser\n",
    "\n",
    "model_kwargs = {\"model\": \"gpt-3.5-turbo\", \"temperature\": 0.3, \"stream\": False}\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Function(DataClass):\n",
    "    name: str = field(metadata={\"desc\": \"The name of the function\"})\n",
    "    args: Dict[str, Any] = field(metadata={\"desc\": \"The arguments of the function\"})\n",
    "\n",
    "\n",
    "generator = Generator(\n",
    "    model_client=ModelClientType.OPENAI(),\n",
    "    model_kwargs=model_kwargs,\n",
    "    template=template,\n",
    "    prompt_kwargs={\n",
    "        # \"tools\": [weather_tool],\n",
    "        \"output_format_str\": YamlOutputParser(Function).format_instructions(),\n",
    "        # \"output_format_str\": Function.to_yaml_signature(),\n",
    "    },\n",
    "    output_processors=YamlOutputParser(Function),\n",
    ")\n",
    "generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt:\n",
      "<SYS>You have these tools available:\n",
      "    <TOOLS>\n",
      "    1. ToolName: _get_current_weather\n",
      "        Tool Description: _get_current_weather(location: str, unit: str = 'fahrenheit')\n",
      "Get the current weather in a given location\n",
      "        Tool Parameters: {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"str\"}, \"unit\": {\"type\": \"str\", \"default\": \"fahrenheit\"}}, \"required\": [\"location\"]}   \n",
      "    __________\n",
      "    </TOOLS>\n",
      "    Your output should be formatted as a standard YAML instance with the following schema:\n",
      "```\n",
      "name: The name of the function (str) (required)\n",
      "args: The arguments of the function (Dict) (required)\n",
      "```\n",
      "\n",
      "-Make sure to always enclose the YAML output in triple backticks (```). Please do not add anything other than valid YAML output!\n",
      "-Follow the YAML formatting conventions with an indent of 2 spaces.\n",
      "-Quote the string values properly.\n",
      "\n",
      "    </SYS>\n",
      "    User: What's the weather like in San Francisco, Tokyo, and Paris in celsius?\n",
      "    You:\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "# check the prompt\n",
    "\n",
    "input_str = \"What's the weather like in San Francisco, Tokyo, and Paris in celsius?\"\n",
    "\n",
    "generator.print_prompt(input_str=input_str, tools=[weather_tool])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Function(name='_get_current_weather', args={'location': 'San Francisco', 'unit': 'celsius'})\n"
     ]
    }
   ],
   "source": [
    "prompt_kwargs = {\n",
    "    \"input_str\": input_str,\n",
    "    \"tools\": [weather_tool],\n",
    "}\n",
    "output = generator(prompt_kwargs=prompt_kwargs)\n",
    "structured_output = Function.from_dict(output.data)\n",
    "print(structured_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": \"celsius\"}\n"
     ]
    }
   ],
   "source": [
    "# call the function\n",
    "\n",
    "function_map = {\"_get_current_weather\": weather_tool}\n",
    "\n",
    "function_name = structured_output.name\n",
    "function_args = structured_output.args\n",
    "function_to_call = function_map[function_name]\n",
    "function_response = function_to_call(**function_args)\n",
    "print(function_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# multiple function calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Generator(\n",
       "  model_kwargs={'model': 'gpt-3.5-turbo', 'temperature': 0.3, 'stream': False}, \n",
       "  (prompt): Prompt(\n",
       "    template: <SYS>You can answer user query with these tools:\n",
       "        <TOOLS>\n",
       "        {% for tool in tools %}\n",
       "        {{ loop.index }}. ToolName: {{ tool.metadata.name }}\n",
       "            Tool Description: {{ tool.metadata.description }}\n",
       "            Tool Parameters: {{ tool.metadata.fn_schema_str }}   \n",
       "        __________\n",
       "        {% endfor %}\n",
       "        </TOOLS>\n",
       "        You can call multiple tools by return a list of the following format:\n",
       "        {{output_format_str}}\n",
       "        </SYS>\n",
       "        User: {{input_str}}\n",
       "        You:\n",
       "        , prompt_kwargs: {'output_format_str': 'Your output should be formatted as a standard YAML instance with the following schema:\\n```\\nname: The name of the function (str) (required)\\nargs: The arguments of the function (Dict) (required)\\n```\\n\\n-Make sure to always enclose the YAML output in triple backticks (```). Please do not add anything other than valid YAML output!\\n-Follow the YAML formatting conventions with an indent of 2 spaces.\\n-Quote the string values properly.\\n'}, prompt_variables: ['input_str', 'output_format_str', 'tools']\n",
       "  )\n",
       "  (model_client): OpenAIClient()\n",
       "  (output_processors): YamlOutputParser(\n",
       "    data_class_for_yaml=<class '__main__.Function'>\n",
       "    (yaml_output_format_prompt): Prompt(\n",
       "      template: Your output should be formatted as a standard YAML instance with the following schema:\n",
       "      ```\n",
       "      {{schema}}\n",
       "      ```\n",
       "      {% if example %}\n",
       "      Here is an example:\n",
       "      ```\n",
       "      {{example}}\n",
       "      ```\n",
       "      {% endif %}\n",
       "      \n",
       "      -Make sure to always enclose the YAML output in triple backticks (```). Please do not add anything other than valid YAML output!\n",
       "      -Follow the YAML formatting conventions with an indent of 2 spaces.\n",
       "      -Quote the string values properly.\n",
       "      , prompt_variables: ['example', 'schema']\n",
       "    )\n",
       "    (output_processors): YamlParser()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "generator = Generator(\n",
    "    model_client=ModelClientType.OPENAI(),\n",
    "    model_kwargs=model_kwargs,\n",
    "    template=multiple_function_call_template,\n",
    "    prompt_kwargs={\n",
    "        # \"tools\": [weather_tool],\n",
    "        \"output_format_str\": YamlOutputParser(Function).format_instructions(),\n",
    "        # \"output_format_str\": Function.to_yaml_signature(),\n",
    "    },\n",
    "    output_processors=YamlOutputParser(Function),\n",
    ")\n",
    "generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GeneratorOutput(data=[{'name': '_get_current_weather', 'args': {'location': 'San Francisco', 'unit': 'celsius'}}, {'name': '_get_current_weather', 'args': {'location': 'Tokyo', 'unit': 'celsius'}}, {'name': '_get_current_weather', 'args': {'location': 'Paris', 'unit': 'celsius'}}], error=None, usage=None, raw_response='```yaml\\n- name: _get_current_weather\\n  args:\\n    location: \"San Francisco\"\\n    unit: \"celsius\"\\n- name: _get_current_weather\\n  args:\\n    location: \"Tokyo\"\\n    unit: \"celsius\"\\n- name: _get_current_weather\\n  args:\\n    location: \"Paris\"\\n    unit: \"celsius\"\\n```')\n",
      "[Function(name='_get_current_weather', args={'location': 'San Francisco', 'unit': 'celsius'}), Function(name='_get_current_weather', args={'location': 'Tokyo', 'unit': 'celsius'}), Function(name='_get_current_weather', args={'location': 'Paris', 'unit': 'celsius'})]\n"
     ]
    }
   ],
   "source": [
    "# run the query\n",
    "\n",
    "output = generator(prompt_kwargs=prompt_kwargs)\n",
    "list_structured_output = [Function.from_dict(item) for item in output.data]\n",
    "print(output)\n",
    "print(list_structured_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": \"celsius\"}\n",
      "{\"location\": \"Tokyo\", \"temperature\": \"10\", \"unit\": \"celsius\"}\n",
      "{\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": \"celsius\"}\n"
     ]
    }
   ],
   "source": [
    "for structured_output in list_structured_output:\n",
    "    function_name = structured_output.name\n",
    "    function_args = structured_output.args\n",
    "    function_to_call = function_map[function_name]\n",
    "    function_response = function_to_call(**function_args)\n",
    "    print(function_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Person(name='John Doe', age=30, address=Address(street='123 Main St', city='Anytown', zipcode='12345'))\n",
      "{'name': 'John Doe', 'age': 30, 'address': {'street': '123 Main St', 'city': 'Anytown', 'zipcode': '12345'}}\n"
     ]
    }
   ],
   "source": [
    "from dataclasses import dataclass, field\n",
    "from typing import Any, Dict\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Address:\n",
    "    street: str\n",
    "    city: str\n",
    "    zipcode: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Person:\n",
    "    name: str\n",
    "    age: int\n",
    "    address: Address\n",
    "\n",
    "\n",
    "# Example instance of the nested dataclasses\n",
    "person = Person(\n",
    "    name=\"John Doe\",\n",
    "    age=30,\n",
    "    address=Address(street=\"123 Main St\", city=\"Anytown\", zipcode=\"12345\"),\n",
    ")\n",
    "print(person)\n",
    "\n",
    "\n",
    "def to_dict(obj: Any) -> Dict[str, Any]:\n",
    "    if hasattr(obj, \"__dataclass_fields__\"):\n",
    "        return {key: to_dict(value) for key, value in obj.__dict__.items()}\n",
    "    elif isinstance(obj, list):\n",
    "        return [to_dict(item) for item in obj]\n",
    "    elif isinstance(obj, dict):\n",
    "        return {key: to_dict(value) for key, value in obj.items()}\n",
    "    else:\n",
    "        return obj\n",
    "\n",
    "\n",
    "# Convert the person instance to a dictionary\n",
    "person_dict = to_dict(person)\n",
    "print(person_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Person(name='John Doe', age=30, addresses=[Address(street='123 Main St', city='Anytown', zipcode='12345'), Address(street='456 Elm St', city='Othertown', zipcode='67890')])\n"
     ]
    }
   ],
   "source": [
    "from typing import List\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Address:\n",
    "    street: str\n",
    "    city: str\n",
    "    zipcode: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Person:\n",
    "    name: str\n",
    "    age: int\n",
    "    addresses: List[Address]\n",
    "\n",
    "\n",
    "# Example instance of the nested dataclasses\n",
    "person = Person(\n",
    "    name=\"John Doe\",\n",
    "    age=30,\n",
    "    addresses=[\n",
    "        Address(street=\"123 Main St\", city=\"Anytown\", zipcode=\"12345\"),\n",
    "        Address(street=\"456 Elm St\", city=\"Othertown\", zipcode=\"67890\"),\n",
    "    ],\n",
    ")\n",
    "print(person)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'name': 'John Doe', 'age': 30, 'addresses': [{'street': '123 Main St', 'city': 'Anytown', 'zipcode': '12345'}, {'street': '456 Elm St', 'city': 'Othertown', 'zipcode': '67890'}]}\n"
     ]
    }
   ],
   "source": [
    "# Convert the person instance to a dictionary\n",
    "person_dict = to_dict(person)\n",
    "print(person_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'data': [{'question': 'What is the capital of France?'}]}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import List, Dict, Optional\n",
    "\n",
    "\n",
    "def dataclass_obj_to_dict(\n",
    "    obj: Any, exclude: Optional[Dict[str, List[str]]] = None, parent_key: str = \"\"\n",
    ") -> Dict[str, Any]:\n",
    "    r\"\"\"Convert a dataclass object to a dictionary.\n",
    "\n",
    "    Supports nested dataclasses, lists, and dictionaries.\n",
    "    Allow exclude keys for each dataclass object.\n",
    "    Example:\n",
    "\n",
    "    .. code-block:: python\n",
    "\n",
    "       from dataclasses import dataclass\n",
    "       from typing import List\n",
    "\n",
    "       @dataclass\n",
    "       class TrecData:\n",
    "           question: str\n",
    "           label: int\n",
    "\n",
    "       @dataclass\n",
    "       class TrecDataList:\n",
    "\n",
    "           data: List[TrecData]\n",
    "           name: str\n",
    "\n",
    "       trec_data = TrecData(question=\"What is the capital of France?\", label=0)\n",
    "       trec_data_list = TrecDataList(data=[trec_data], name=\"trec_data_list\")\n",
    "\n",
    "       dataclass_obj_to_dict(trec_data_list, exclude={\"TrecData\": [\"label\"], \"TrecDataList\": [\"name\"]})\n",
    "\n",
    "       # Output:\n",
    "       # {'data': [{'question': 'What is the capital of France?'}], 'name': 'trec_data_list'}\n",
    "\n",
    "    \"\"\"\n",
    "    if exclude is None:\n",
    "        exclude = {}\n",
    "\n",
    "    obj_class_name = obj.__class__.__name__\n",
    "    current_exclude = exclude.get(obj_class_name, [])\n",
    "\n",
    "    if hasattr(obj, \"__dataclass_fields__\"):\n",
    "        return {\n",
    "            key: dataclass_obj_to_dict(value, exclude, parent_key=key)\n",
    "            for key, value in obj.__dict__.items()\n",
    "            if key not in current_exclude\n",
    "        }\n",
    "    elif isinstance(obj, list):\n",
    "        return [dataclass_obj_to_dict(item, exclude, parent_key) for item in obj]\n",
    "    elif isinstance(obj, dict):\n",
    "        return {\n",
    "            key: dataclass_obj_to_dict(value, exclude, parent_key)\n",
    "            for key, value in obj.items()\n",
    "        }\n",
    "    else:\n",
    "        return obj\n",
    "\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import List\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TrecData:\n",
    "    question: str\n",
    "    label: int\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TrecDataList:\n",
    "\n",
    "    data: List[TrecData]\n",
    "    name: str\n",
    "\n",
    "\n",
    "trec_data = TrecData(question=\"What is the capital of France?\", label=0)\n",
    "trec_data_list = TrecDataList(data=[trec_data], name=\"trec_data_list\")\n",
    "\n",
    "dataclass_obj_to_dict(\n",
    "    trec_data_list, exclude={\"TrecData\": [\"label\"], \"TrecDataList\": [\"name\"]}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Type\n",
    "\n",
    "\n",
    "def dataclass_obj_from_dict(cls: Type[Any], data: Dict[str, Any]) -> Any:\n",
    "    if hasattr(cls, \"__dataclass_fields__\"):\n",
    "        fieldtypes = {f.name: f.type for f in cls.__dataclass_fields__.values()}\n",
    "        return cls(\n",
    "            **{\n",
    "                key: dataclass_obj_from_dict(fieldtypes[key], value)\n",
    "                for key, value in data.items()\n",
    "            }\n",
    "        )\n",
    "    elif isinstance(data, list):\n",
    "        return [dataclass_obj_from_dict(cls.__args__[0], item) for item in data]\n",
    "    elif isinstance(data, dict):\n",
    "        return {\n",
    "            key: dataclass_obj_from_dict(cls.__args__[1], value)\n",
    "            for key, value in data.items()\n",
    "        }\n",
    "    else:\n",
    "        return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrecDataList(data=[TrecData(question='What is the capital of France?', label=0)], name='trec_data_list')"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataclass_obj_from_dict(TrecDataList, dataclass_obj_to_dict(trec_data_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "TrecData.__init__() missing 1 required positional argument: 'label'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[25], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mdataclass_obj_from_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mTrecDataList\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataclass_obj_to_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrec_data_list\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexclude\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mTrecData\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlabel\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mTrecDataList\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mname\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[23], line 5\u001b[0m, in \u001b[0;36mdataclass_obj_from_dict\u001b[0;34m(cls, data)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(\u001b[38;5;28mcls\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__dataclass_fields__\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m      4\u001b[0m     fieldtypes \u001b[38;5;241m=\u001b[39m {f\u001b[38;5;241m.\u001b[39mname: f\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__dataclass_fields__\u001b[38;5;241m.\u001b[39mvalues()}\n\u001b[0;32m----> 5\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mcls\u001b[39m(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[43m{\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataclass_obj_from_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfieldtypes\u001b[49m\u001b[43m[\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m}\u001b[49m)\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, \u001b[38;5;28mlist\u001b[39m):\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [dataclass_obj_from_dict(\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__args__[\u001b[38;5;241m0\u001b[39m], item) \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m data]\n",
      "Cell \u001b[0;32mIn[23], line 5\u001b[0m, in \u001b[0;36m<dictcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(\u001b[38;5;28mcls\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__dataclass_fields__\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m      4\u001b[0m     fieldtypes \u001b[38;5;241m=\u001b[39m {f\u001b[38;5;241m.\u001b[39mname: f\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__dataclass_fields__\u001b[38;5;241m.\u001b[39mvalues()}\n\u001b[0;32m----> 5\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mcls\u001b[39m(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m{key: \u001b[43mdataclass_obj_from_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfieldtypes\u001b[49m\u001b[43m[\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m data\u001b[38;5;241m.\u001b[39mitems()})\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, \u001b[38;5;28mlist\u001b[39m):\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [dataclass_obj_from_dict(\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__args__[\u001b[38;5;241m0\u001b[39m], item) \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m data]\n",
      "Cell \u001b[0;32mIn[23], line 7\u001b[0m, in \u001b[0;36mdataclass_obj_from_dict\u001b[0;34m(cls, data)\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mcls\u001b[39m(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m{key: dataclass_obj_from_dict(fieldtypes[key], value) \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m data\u001b[38;5;241m.\u001b[39mitems()})\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, \u001b[38;5;28mlist\u001b[39m):\n\u001b[0;32m----> 7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m[\u001b[49m\u001b[43mdataclass_obj_from_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__args__\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mitem\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, \u001b[38;5;28mdict\u001b[39m):\n\u001b[1;32m      9\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m {key: dataclass_obj_from_dict(\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__args__[\u001b[38;5;241m1\u001b[39m], value) \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m data\u001b[38;5;241m.\u001b[39mitems()}\n",
      "Cell \u001b[0;32mIn[23], line 7\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mcls\u001b[39m(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m{key: dataclass_obj_from_dict(fieldtypes[key], value) \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m data\u001b[38;5;241m.\u001b[39mitems()})\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, \u001b[38;5;28mlist\u001b[39m):\n\u001b[0;32m----> 7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[43mdataclass_obj_from_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__args__\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m data]\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, \u001b[38;5;28mdict\u001b[39m):\n\u001b[1;32m      9\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m {key: dataclass_obj_from_dict(\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__args__[\u001b[38;5;241m1\u001b[39m], value) \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m data\u001b[38;5;241m.\u001b[39mitems()}\n",
      "Cell \u001b[0;32mIn[23], line 5\u001b[0m, in \u001b[0;36mdataclass_obj_from_dict\u001b[0;34m(cls, data)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(\u001b[38;5;28mcls\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__dataclass_fields__\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m      4\u001b[0m     fieldtypes \u001b[38;5;241m=\u001b[39m {f\u001b[38;5;241m.\u001b[39mname: f\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__dataclass_fields__\u001b[38;5;241m.\u001b[39mvalues()}\n\u001b[0;32m----> 5\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m{\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataclass_obj_from_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfieldtypes\u001b[49m\u001b[43m[\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, \u001b[38;5;28mlist\u001b[39m):\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [dataclass_obj_from_dict(\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m__args__[\u001b[38;5;241m0\u001b[39m], item) \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m data]\n",
      "\u001b[0;31mTypeError\u001b[0m: TrecData.__init__() missing 1 required positional argument: 'label'"
     ]
    }
   ],
   "source": [
    "dataclass_obj_from_dict(\n",
    "    TrecDataList,\n",
    "    dataclass_obj_to_dict(\n",
    "        trec_data_list, exclude={\"TrecData\": [\"label\"], \"TrecDataList\": [\"name\"]}\n",
    "    ),\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lightrag-project",
   "language": "python",
   "name": "light-rag-project"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
